如何在矩阵乘法任务中避免循环以加速PyTorch代码运行？

Great question! Ditching loops for vectorized operations in PyTorch is one of the best ways to speed up your code—especially when working with large batches—since it lets you leverage GPU acceleration and avoid the overhead of Python-level loops. Let's walk through how to optimize your code step by step.

Key Issues with the Original Code #

Before diving into the fix, let's note what's slowing things down:

• Python loop: Iterating over 1000 samples one by one means you're missing out on PyTorch's parallel computing capabilities.
• Numpy-Torch conversions: Switching between numpy arrays and torch tensors creates unnecessary data copy overhead.
• Per-sample operations: Generating theta/alpha and constructing vectors one at a time is inefficient compared to batch processing.

Optimized Vectorized Code #

Here's the fully optimized version, with explanations for each step:

import torch
import numpy as np

# Define parameters clearly
batch_size = 1000
K = 10
theta_min = 0.0
theta_max = np.pi
sigma = 10.0

# Dummy neural network output (same as your original)
y = torch.ones((batch_size, K))

# --------------------------
# Step 1: Batch generate random variables with PyTorch
# --------------------------
# Generate 1000 theta values in one go (shape: (1000,))
theta = torch.rand(batch_size) * (theta_max - theta_min) + theta_min
# Generate 1000 alpha values (shape: (1000,))
alpha = sigma * torch.randn(batch_size)

# --------------------------
# Step 2: Batch construct complex vectors
# --------------------------
# Create a tensor of [0, 1, ..., K-1] (shape: (K,))
arange = torch.arange(K, dtype=torch.float32)
# Compute sin(theta) for all samples (shape: (1000,))
sin_theta = torch.sin(theta)

# Broadcast calculations to create terms for all samples (shape: (1000, K))
# arange[None, :] adds a batch dimension, sin_theta[:, None] adds a feature dimension
terms = -1j * arange[None, :] * np.pi * sin_theta[:, None]
# Exponentiate and add a trailing dimension to make it (1000, K, 1)
vector_batch = torch.exp(terms).unsqueeze(-1)

# --------------------------
# Step 3: Batch matrix multiplication
# --------------------------
# Compute vector @ vector.T for all samples (shape: (1000, K, K))
matrix_batch = vector_batch @ vector_batch.transpose(1, 2)
# Reshape y to (1000, K, 1) for batch matrix multiplication
y_batch = y.unsqueeze(-1)

# Perform all matrix multiplications in parallel
temp = matrix_batch @ y_batch  # Shape: (1000, K, 1)
# Multiply by alpha (broadcast alpha to (1000, 1, 1) to match dimensions)
z = alpha[:, None, None] * temp
# Remove the trailing dimension to get back (1000, K)
z = z.squeeze(-1)

Why This Is Faster #

• No loops: All operations are performed on batches of data, letting PyTorch parallelize computations across CPU cores or GPU.
• No data conversions: Everything stays in PyTorch tensors, eliminating the overhead of copying data between numpy and torch.
• Broadcasted operations: We use PyTorch's broadcasting rules to avoid explicit repetition of tensors, keeping memory usage efficient.

Bonus: GPU Acceleration #

If you have access to a GPU, simply move all tensors to the device to get even faster results:

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
y = y.to(device)
theta = theta.to(device)
alpha = alpha.to(device)
arange = arange.to(device)

Verify Correctness #

To ensure the optimized code matches your original, set a random seed for both numpy and PyTorch, then compare a few samples:

# Set seeds for reproducibility
np.random.seed(42)
torch.manual_seed(42)

# Run original loop code for a few samples and compare with z from the optimized code

内容的提问来源于stack exchange，提问作者Josemi