PyTorch Advanced Tensor Operations
Introduction
In this tutorial, we'll dive into advanced tensor operations in PyTorch. Once you've mastered the basics of creating and manipulating tensors, these advanced operations will help you write more efficient and expressive code for your deep learning projects. We'll cover broadcasting, reshaping, advanced indexing, masking operations, and more - all essential techniques for efficient tensor manipulation.
Broadcasting in PyTorch
Broadcasting is a powerful mechanism that allows PyTorch to work with tensors of different shapes when performing arithmetic operations. Instead of explicitly reshaping tensors, PyTorch automatically "broadcasts" the smaller tensor to match the shape of the larger one.
How Broadcasting Works
To understand broadcasting, remember these rules:
- Each tensor must have at least one dimension
- Dimensions are compared from right to left
- Each dimension must either be equal, or one of them must be 1, or one doesn't exist
Let's look at some examples:
import torch
# Creating tensors of different shapes
a = torch.tensor([1, 2, 3]) # Shape: [3]
b = torch.tensor([[1], [2]]) # Shape: [2,1]
# Broadcasting in action
c = a + b # b will be broadcast to match a's dimensions
print(f"Tensor a shape: {a.shape}")
print(f"Tensor b shape: {b.shape}")
print(f"Result c shape: {c.shape}")
print(c)
Output:
Tensor a shape: torch.Size([3])
Tensor b shape: torch.Size([2, 1])
Result c shape: torch.Size([2, 3])
tensor([[2, 3, 4],
[3, 4, 5]])
In this example, the shapes [3] and [2, 1] are broadcast to [2, 3]. The tensor a is expanded to have 2 rows, and the tensor b is expanded to have 3 columns.
Practical Use Case: Batch Processing
Broadcasting is particularly useful for batch processing in deep learning:
# Weight matrix for a neural network layer
weights = torch.randn(128, 10) # 128 features, 10 outputs
# Batch of inputs
batch = torch.randn(32, 128) # 32 samples, 128 features
# Process the whole batch at once with broadcasting
outputs = torch.matmul(batch, weights) # Shape: [32, 10]
print(f"Weights shape: {weights.shape}")
print(f"Batch shape: {batch.shape}")
print(f"Outputs shape: {outputs.shape}")
Output:
Weights shape: torch.Size([128, 10])
Batch shape: torch.Size([32, 128])
Outputs shape: torch.Size([32, 10])
Advanced Reshaping Operations
PyTorch offers several methods to change the shape of tensors while preserving the data.
View, Reshape, and Transpose
import torch
# Create a tensor
x = torch.arange(12)
print(f"Original tensor: {x}")
print(f"Shape: {x.shape}")
# Using view
x_viewed = x.view(3, 4)
print(f"\nAfter view(3, 4):\n{x_viewed}")
# Using reshape (more flexible than view)
x_reshaped = x.reshape(2, 6)
print(f"\nAfter reshape(2, 6):\n{x_reshaped}")
# Transpose - swap dimensions
x_transposed = x_reshaped.transpose(0, 1)
print(f"\nAfter transpose(0, 1):\n{x_transposed}")
print(f"Transposed shape: {x_transposed.shape}")
# Permute - generalized transpose
x_3d = x.reshape(2, 2, 3)
print(f"\n3D tensor:\n{x_3d}")
x_permuted = x_3d.permute(2, 0, 1)
print(f"\nAfter permute(2, 0, 1):\n{x_permuted}")
print(f"Permuted shape: {x_permuted.shape}")
Output:
Original tensor: tensor([ 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11])
Shape: torch.Size([12])
After view(3, 4):
tensor([[ 0, 1, 2, 3],
[ 4, 5, 6, 7],
[ 8, 9, 10, 11]])
After reshape(2, 6):
tensor([[ 0, 1, 2, 3, 4, 5],
[ 6, 7, 8, 9, 10, 11]])
After transpose(0, 1):
tensor([[ 0, 6],
[ 1, 7],
[ 2, 8],
[ 3, 9],
[ 4, 10],
[ 5, 11]])
Transposed shape: torch.Size([6, 2])
3D tensor:
tensor([[[ 0, 1, 2],
[ 3, 4, 5]],
[[ 6, 7, 8],
[ 9, 10, 11]]])
After permute(2, 0, 1):
tensor([[[ 0, 3],
[ 6, 9]],
[[ 1, 4],
[ 7, 10]],
[[ 2, 5],
[ 8, 11]]])
Permuted shape: torch.Size([3, 2, 2])
View vs Reshape
The key difference:
view()returns a tensor that shares storage with the original tensorreshape()might return a copy if the tensor is not contiguous
# Create a tensor and transpose it
a = torch.arange(9).reshape(3, 3)
b = a.t() # transpose
# This will fail because b is not contiguous
try:
c = b.view(9)
except Exception as e:
print(f"Error with view: {e}")
# This works because reshape handles non-contiguous tensors
d = b.reshape(9)
print(f"Reshape result: {d}")
# Check if tensors are contiguous
print(f"Is a contiguous? {a.is_contiguous()}")
print(f"Is b contiguous? {b.is_contiguous()}")
Output:
Error with view: view size is not compatible with input tensor's size and stride (at least one dimension spans across two contiguous subspaces). Use .reshape(...) instead.
Reshape result: tensor([0, 3, 6, 1, 4, 7, 2, 5, 8])
Is a contiguous? True
Is b contiguous? False
Advanced Indexing and Slicing
PyTorch supports advanced indexing operations similar to NumPy.
Boolean Masking
import torch
# Create a tensor
x = torch.randn(5, 5)
print(f"Original tensor:\n{x}")
# Create a boolean mask
mask = x > 0
print(f"\nBoolean mask (x > 0):\n{mask}")
# Apply the mask
positive_values = x[mask]
print(f"\nPositive values only:\n{positive_values}")
# Count positive values
print(f"Number of positive values: {mask.sum().item()}")
# Replace negative values with zeros
x_no_neg = x.clone()
x_no_neg[x < 0] = 0
print(f"\nAfter replacing negative values with zeros:\n{x_no_neg}")
Fancy Indexing
# Create a tensor
x = torch.arange(10, 20)
print(f"Original tensor: {x}")
# Select specific indices
indices = torch.tensor([0, 2, 4, 8])
selected = x[indices]
print(f"\nSelected elements at indices {indices}: {selected}")
# For 2D tensors
y = torch.arange(25).reshape(5, 5)
print(f"\n2D tensor:\n{y}")
# Select specific rows
row_indices = torch.tensor([0, 2, 4])
selected_rows = y[row_indices]
print(f"\nSelected rows:\n{selected_rows}")
# Select specific elements with row and column indices
row_indices = torch.tensor([0, 1, 2])
col_indices = torch.tensor([1, 3, 4])
elements = y[row_indices, col_indices]
print(f"\nSelected elements: {elements}")
Advanced Mathematical Operations
PyTorch offers a wide range of mathematical functions for tensors.
Matrix Operations
import torch
# Create tensors
a = torch.randn(3, 4)
b = torch.randn(4, 5)
# Matrix multiplication
c = torch.matmul(a, b)
print(f"Matrix multiplication result shape: {c.shape}")
# Another way to perform matrix multiplication
c_alt = a @ b
print(f"Using @ operator shape: {c_alt.shape}")
print(f"Are they the same? {torch.allclose(c, c_alt)}")
# Batched matrix multiplication
batch_a = torch.randn(10, 3, 4) # 10 batches of 3x4 matrices
batch_b = torch.randn(10, 4, 5) # 10 batches of 4x5 matrices
batch_c = torch.bmm(batch_a, batch_b) # Batched matrix multiply
print(f"Batched matrix multiplication result shape: {batch_c.shape}")
# Computing the inverse of a matrix
square_matrix = torch.randn(4, 4)
try:
inverse = torch.inverse(square_matrix)
print(f"Matrix inverse computed successfully")
# Verify: A * A^(-1) = I
identity = torch.matmul(square_matrix, inverse)
print(f"Verification (should be close to identity matrix):\n{identity}")
except Exception as e:
print(f"Error computing inverse: {e}")
Eigenvalues and Eigenvectors
# Create a symmetric matrix
symmetric = torch.randn(4, 4)
symmetric = symmetric + symmetric.t() # Make it symmetric
print(f"Symmetric matrix:\n{symmetric}")
# Compute eigenvalues and eigenvectors
eigenvalues, eigenvectors = torch.linalg.eigh(symmetric)
print(f"\nEigenvalues: {eigenvalues}")
print(f"Eigenvectors shape: {eigenvectors.shape}")
# Verify Av = λv
for i in range(len(eigenvalues)):
v = eigenvectors[:, i]
Av = torch.matmul(symmetric, v)
lambda_v = eigenvalues[i] * v
print(f"Verification for eigenvalue {i}:")
print(f"Av ≈ λv: {torch.allclose(Av, lambda_v, atol=1e-5)}")
Tensor Memory Management and Efficiency
Efficient memory management is crucial when working with large tensors.
In-place Operations
In-place operations modify a tensor directly without allocating new memory.
import torch
# Create a tensor
x = torch.ones(3, 3)
print(f"Original x:\n{x}")
# Regular operation (creates new tensor)
y = x + 2
print(f"y = x + 2:\n{y}")
print(f"Original x is unchanged:\n{x}")
# In-place operation (modifies original tensor)
x.add_(2) # Note the underscore suffix indicating in-place
print(f"After x.add_(2):\n{x}")
Output:
Original x:
tensor([[1., 1., 1.],
[1., 1., 1.],
[1., 1., 1.]])
y = x + 2:
tensor([[3., 3., 3.],
[3., 3., 3.],
[3., 3., 3.]])
Original x is unchanged:
tensor([[1., 1., 1.],
[1., 1., 1.],
[1., 1., 1.]])
After x.add_(2):
tensor([[3., 3., 3.],
[3., 3., 3.],
[3., 3., 3.]])
Sharing Memory Between Tensors
Understanding how tensors share memory can help you optimize your code.
# Create a tensor
x = torch.ones(5, 5)
# Create a view (shares storage)
y = x.view(25)
print(f"x shape: {x.shape}, y shape: {y.shape}")
# Modify y
y[0] = 100
# This affects x as well
print(f"Modified first element of y to 100")
print(f"x[0,0] = {x[0,0]}") # Will show 100
# Create a copy (independent)
z = x.clone()
z[0, 0] = 200
print(f"Modified first element of z to 200")
print(f"x[0,0] = {x[0,0]}") # Still 100, not affected
Moving Tensors Between Devices
For performance optimization, especially in deep learning, you often need to move tensors between CPU and GPU.
# Create a tensor on CPU
x_cpu = torch.randn(3, 3)
print(f"x_cpu device: {x_cpu.device}")
# Check if GPU is available
if torch.cuda.is_available():
# Move to GPU
x_gpu = x_cpu.cuda() # or x_cpu.to("cuda")
print(f"x_gpu device: {x_gpu.device}")
# Move back to CPU
x_back_to_cpu = x_gpu.cpu()
print(f"x_back_to_cpu device: {x_back_to_cpu.device}")
# Operations between tensors must be on the same device
try:
result = x_cpu + x_gpu # This will fail
except Exception as e:
print(f"Error: {e}")
# Correct way to add tensors from different devices
result = x_cpu + x_gpu.cpu() # Move to same device first
print(f"Addition successful after moving to same device")
else:
print("CUDA not available on this system")
Real-World Applications
Let's explore some practical applications of advanced tensor operations.
Image Processing with Tensors
import torch
import matplotlib.pyplot as plt
import numpy as np
# Create a simple 8x8 image tensor
image = torch.zeros(8, 8)
image[2:6, 2:6] = 1.0 # Create a white square
print(f"Original image:\n{image}")
# Apply a convolution for edge detection
kernel = torch.tensor([[-1, -1, -1],
[-1, 8, -1],
[-1, -1, -1]], dtype=torch.float32)
# Reshape for convolution
image_3d = image.view(1, 1, 8, 8) # Add batch and channel dimensions
kernel_4d = kernel.view(1, 1, 3, 3) # Add output and input channel dimensions
# Perform convolution
edge_detected = torch.nn.functional.conv2d(image_3d, kernel_4d, padding=1)
# Remove extra dimensions
edge_detected = edge_detected.squeeze()
print(f"\nEdge detected image:\n{edge_detected}")
# Visualize (this code would work in a notebook or with plt.show())
"""
plt.figure(figsize=(10, 4))
plt.subplot(1, 2, 1)
plt.imshow(image.numpy(), cmap='gray')
plt.title("Original Image")
plt.subplot(1, 2, 2)
plt.imshow(edge_detected.numpy(), cmap='gray')
plt.title("Edge Detected")
plt.tight_layout()
"""
Natural Language Processing Example
import torch
# Let's create a simple one-hot encoding for a vocabulary
vocab = ["hello", "world", "pytorch", "tensor", "operations"]
# Create a one-hot encoding function
def one_hot_encode(word, vocab):
# Create a tensor of zeros
encoding = torch.zeros(len(vocab))
# Set the position corresponding to the word to 1
if word in vocab:
encoding[vocab.index(word)] = 1
return encoding
# Encode some words
words = ["hello", "pytorch", "operations"]
encodings = torch.stack([one_hot_encode(word, vocab) for word in words])
print(f"One-hot encodings:\n{encodings}")
# Create a simple embedding matrix (in real applications, this would be learned)
embedding_dim = 3
embedding_matrix = torch.randn(len(vocab), embedding_dim)
print(f"\nEmbedding matrix ({len(vocab)}x{embedding_dim}):\n{embedding_matrix}")
# Convert one-hot encodings to dense embeddings
dense_embeddings = torch.matmul(encodings, embedding_matrix)
print(f"\nDense embeddings:\n{dense_embeddings}")
# For each word in our list, show its embedding
for i, word in enumerate(words):
print(f"Embedding for '{word}': {dense_embeddings[i]}")
Summary
In this tutorial, we've covered advanced tensor operations in PyTorch:
- Broadcasting - Efficiently works with tensors of different shapes
- Reshaping Operations - View, reshape, transpose, and permute
- Advanced Indexing - Boolean masking and fancy indexing
- Matrix Operations - Matrix multiplication, inverses, eigenvalues
- Memory Management - In-place operations, sharing memory, device management
- Real-World Applications - Image processing and NLP examples
These advanced operations form the foundation for efficient tensor manipulation in deep learning projects. By mastering these techniques, you'll be able to write cleaner, faster, and more memory-efficient PyTorch code.
Additional Resources and Exercises
Resources
- PyTorch Documentation on Tensor Operations
- PyTorch Broadcasting Semantics
- PyTorch Memory Management Documentation
Exercises
-
Broadcasting Practice: Create two tensors of shapes
[3,1,4]and[1,5,4]and add them together. What will be the shape of the result? -
Advanced Indexing: Create a 5x5 tensor with random values. Extract all elements where the value is greater than the mean of the tensor.
-
Matrix Operations: Create a batched matrix multiplication scenario for a simple neural network layer with batch size 32, input features 64, and output features 10.
-
Memory Efficiency: Modify a tensor in-place using five different operations. Verify that each operation didn't create a new tensor.
-
Real-World Application: Implement a simple image blurring operation using convolution with a Gaussian kernel on a tensor representing an image.
Good luck with your PyTorch journey!
If you spot any mistakes on this website, please let me know at [email protected]. I’d greatly appreciate your feedback! :)