Back to blog
• 5 min read • Deep Learning

PyTorch Basics: Tensors and Autograd

PyTorch has become the framework of choice for researchers and engineers alike, offering an intuitive Pythonic interface while maintaining the performance needed for production workloads. This guide covers the foundational concepts you need to start building with PyTorch.

What Makes PyTorch Special?

PyTorch stands out from other deep learning frameworks due to its dynamic computation graph. Unlike static graph frameworks where you define the entire computation upfront, PyTorch builds the graph on-the-fly as operations execute. This means:

  • Intuitive debugging: You can use standard Python debugging tools like pdb
  • Dynamic architectures: Build models that change structure based on input
  • Pythonic code: Write natural Python instead of framework-specific DSLs

PyTorch provides tensor computation with GPU acceleration and deep neural networks built on a tape-based autograd system.

Understanding Tensors

Tensors are the fundamental data structure in PyTorch—multi-dimensional arrays that can reside on CPUs or GPUs. Think of them as NumPy arrays with superpowers.

Creating and Manipulating Tensors

import torch

# Creating tensors from data
x = torch.tensor([1, 2, 3])
y = torch.ones(3, 4)
z = torch.zeros(2, 3, 4)

# Tensor with random values
a = torch.randn(3, 3)
b = torch.randn(3, 3)

# Basic operations
c = torch.mm(a, b)  # Matrix multiplication
d = a + b           # Element-wise addition
e = a * b           # Element-wise multiplication

print(f"Shape of a: {a.shape}")
print(f"Device: {a.device}")

Moving Tensors to GPU

One of PyTorch’s strengths is seamless GPU acceleration:

# Check if GPU is available
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# Create tensor on GPU
tensor_gpu = torch.randn(3, 3, device=device)

# Or move existing tensor
x = torch.randn(3, 3)
x_gpu = x.to(device)

Automatic Differentiation with Autograd

The magic of PyTorch lies in its automatic differentiation engine, autograd. It tracks operations on tensors and automatically computes gradients—a process essential for training neural networks.

How Autograd Works

When you create a tensor with requires_grad=True, PyTorch starts tracking all operations on that tensor:

# Create tensors with gradient tracking
x = torch.tensor(2.0, requires_grad=True)
y = torch.tensor(3.0, requires_grad=True)

# Define computation
z = x**2 + y**3

# Compute gradients
z.backward()

print(f"dz/dx: {x.grad}")  # 4.0 (2*2)
print(f"dz/dy: {y.grad}")  # 27.0 (3*9)

The backward() call computes gradients using the chain rule, populating the .grad attribute of each tensor that requires gradients.

Complex Computation Graphs

Autograd handles complex graphs with ease:

x = torch.tensor([1.0, 2.0, 3.0], requires_grad=True)

# Multiple operations
y = x * 2
z = y.sum()
w = z ** 2

# Backward pass
w.backward()

# Gradient flows through entire computation
print(f"dw/dx: {x.grad}")

Building Neural Networks

PyTorch provides the torch.nn module for building neural network architectures. Here’s a complete example:

import torch.nn as nn
import torch.nn.functional as F

class NeuralNetwork(nn.Module):
    def __init__(self, input_size, hidden_size, num_classes):
        super(NeuralNetwork, self).__init__()
        self.layer1 = nn.Linear(input_size, hidden_size)
        self.relu = nn.ReLU()
        self.dropout = nn.Dropout(0.2)
        self.layer2 = nn.Linear(hidden_size, num_classes)

    def forward(self, x):
        out = self.layer1(x)
        out = self.relu(out)
        out = self.dropout(out)
        out = self.layer2(out)
        return out

# Create model
model = NeuralNetwork(784, 256, 10)

# Move to GPU if available
model = model.to(device)

Training Loop

A typical training loop involves:

# Define loss and optimizer
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)

# Training loop
for epoch in range(num_epochs):
    for batch_idx, (data, target) in enumerate(train_loader):
        # Move data to device
        data, target = data.to(device), target.to(device)

        # Forward pass
        output = model(data)
        loss = criterion(output, target)

        # Backward pass
        optimizer.zero_grad()  # Clear previous gradients
        loss.backward()        # Compute gradients
        optimizer.step()       # Update weights

        if batch_idx % 100 == 0:
            print(f'Epoch: {epoch}, Batch: {batch_idx}, Loss: {loss.item():.4f}')

Best Practices

Always call optimizer.zero_grad(): Gradients accumulate by default. Forgetting to zero them leads to incorrect updates.

Use torch.no_grad() for inference: This disables gradient computation, saving memory and speeding up evaluation:

model.eval()
with torch.no_grad():
    predictions = model(test_data)

Monitor GPU memory: PyTorch doesn’t automatically release GPU memory. Use torch.cuda.empty_cache() if needed.

Save and load models properly:

# Save
torch.save(model.state_dict(), 'model.pth')

# Load
model = NeuralNetwork(784, 256, 10)
model.load_state_dict(torch.load('model.pth'))
model.eval()

Common Pitfalls

  1. Modifying tensors in-place: Operations like x += 1 can break the computation graph. Use x = x + 1 instead.

  2. Forgetting to move data to GPU: Always check that both model and data are on the same device.

  3. Not calling model.train() or model.eval(): Dropout and batch normalization behave differently during training and evaluation.

Conclusion

PyTorch’s dynamic computation graph and intuitive API make it an excellent choice for both research and production. Mastering tensors and autograd provides the foundation for building complex deep learning systems. The skills learned here transfer directly to production deployment with TorchScript or ONNX export.

For production environments, PyTorch offers TorchServe for model serving and integrates seamlessly with MLflow and other MLOps tools. The knowledge of PyTorch fundamentals serves as a springboard into the broader deep learning ecosystem.

Back to top