PyTorch空间变换核心：AffineGrid与GridSample原理与实战

1. 这两个函数到底在干啥：从图像变形的底层逻辑讲起

AffineGrid 和 GridSample 是 PyTorch 图像空间变换中真正扛大梁的一对组合，它们不直接处理像素值，而是联手完成“坐标重映射”这件事——这恰恰是所有几何变换（旋转、缩放、平移、仿射扭曲、透视变形）最核心的数学本质。很多人一看到这两个函数名就头大，觉得是高级API、是底层黑盒，其实拆开来看，它们分工极其清晰：AffineGrid 负责“画地图”，GridSample 负责“按图索骥”。你给它一个仿射变换矩阵，AffineGrid 就能算出目标图像上每个像素点，在原始图像里该去哪找颜色；而 GridSample 拿着这张“坐标地图”，真刀真枪地把原始图像的像素值采样、插值、填到目标位置上。整个过程完全可导，所以它天然支持反向传播，是训练可变形卷积、空间变换网络（STN）、风格迁移中形变模块、医学图像配准模型的基石。

我第一次在某高校实验室复现一篇 STN 论文时，卡在了这里整整三天。不是代码报错，而是输出图像全是模糊的灰块，或者出现诡异的条纹撕裂。后来才发现，问题根本不在模型结构，而在 AffineGrid 生成的网格本身——它默认输出的是归一化坐标（-1 到 1），而我当时误以为是像素坐标（0 到 H-1），直接拿去做了整数索引，结果全成了越界访问。这个坑，几乎每个刚接触空间变换的同学都会踩。它不像torch.nn.functional.interpolate那样“傻瓜式”调用，它要求你对坐标系、采样原理、插值边界有明确的物理直觉。比如，为什么 AffineGrid 的输入是 2x3 矩阵而不是 3x3？因为 PyTorch 的仿射变换只处理二维平面，第三行 [0,0,1] 是齐次坐标的固定项，被框架隐式补全了，你只需提供前两行，即控制 x 和 y 方向的线性组合与平移项。再比如，GridSample 的mode参数选bilinear还是nearest，不只是影响画质，更直接影响梯度流的平滑性——nearest模式下，梯度在像素边界处是不连续的，某些对梯度敏感的损失函数（如感知损失）会因此训练不稳定。这些细节，文档里不会明说，但实操中处处是雷区。

这两个函数之所以常被并称“兄弟”，是因为它们极少单独出现。你几乎找不到一个只用 AffineGrid 不用 GridSample 的合理场景，反之亦然。它们构成了一条完整的、端到端可微的“坐标变换流水线”。理解它们，不是为了多记两个函数名，而是为了掌握 PyTorch 处理空间几何操作的通用范式：先定义变换规则（AffineGrid），再执行采样动作（GridSample）。这种“声明式定义 + 命令式执行”的分离设计，让模型具备了极强的灵活性。你可以用同一个 GridSample，喂给它由 AffineGrid 生成的仿射网格，也可以喂给由自定义神经网络预测出的稠密光流场，甚至可以喂给手写的正弦波扰动网格来实现动态纹理扭曲。这种解耦，正是现代深度学习框架强大表达力的体现。所以，别把它当成两个孤立的工具，而要当成一把打开空间操作大门的钥匙——门后是 STN、Deformable Conv、Thin Plate Spline、Neural Radiance Fields（NeRF）中的相机姿态变换……几乎所有需要“动”图像位置的地方，都有它们的身影。

2. 核心原理拆解：从数学公式到 PyTorch 实现的逐层穿透

2.1 AffineGrid：一张动态生成的“寻址地图”

AffineGrid 的核心任务，是根据一个或多个 2x3 的仿射变换矩阵，为指定尺寸的目标特征图（通常是 batch x C x H x W）生成一个对应的采样网格（grid）。这个网格的形状是batch x H x W x 2，其中最后的2维度分别代表该位置在源图像中应采样的x 坐标和y 坐标。关键在于，这个坐标不是以像素为单位的整数，而是归一化的浮点坐标，范围严格限定在[-1, 1]区间内。这个设计是 PyTorch 与 OpenGL、CUDA 纹理采样硬件高度对齐的结果，也是其高性能的底层保障。

我们来推导一下这个归一化坐标的来源。假设你有一个标准的 2D 仿射变换矩阵A，其形式为：

A = [[a11, a12, t_x], [a21, a22, t_y]]

其中t_x,t_y是平移分量。对于目标图像上一个位于(i, j)的像素（i为行索引，j为列索引），其在源图像中的理论坐标(x_src, y_src)应为：

[x_src] [a11 a12 t_x] [j] [y_src] = [a21 a22 t_y] [i] [ 1 ] [ 0 0 1 ] [1]

注意，这里j对应 x 轴（列），i对应 y 轴（行），这是图像坐标系的惯例。但 AffineGrid 并不直接输出(x_src, y_src)，而是将其映射到[-1, 1]区间。映射规则是：源图像的左上角(0, 0)映射为(-1, -1)，右下角(W-1, H-1)映射为(1, 1)。因此，最终的归一化坐标(x_norm, y_norm)为：

x_norm = 2 * (x_src / (W - 1)) - 1 y_norm = 2 * (y_src / (H - 1)) - 1

这就是 AffineGrid 内部执行的完整计算。它本质上是一个批量的、向量化的坐标变换器。当你传入一个batch_size=4的theta张量（shape:4x2x3），它会一次性为这 4 个样本，各自生成一张HxWx2的网格，总输出 shape 为4xHxWx2。这个过程没有循环，全部由底层 CUDA kernel 高效完成，这也是它比手动写 for 循环快几个数量级的原因。

提示：theta矩阵的排列顺序极易混淆。PyTorch 要求theta的 shape 是(N, 2, 3)，其中N是 batch size。第一行[a11, a12, t_x]控制 x 坐标的变换，第二行[a21, a22, t_y]控制 y 坐标的变换。如果你用 OpenCV 的getAffineTransform得到一个2x3矩阵，它通常就是正确的格式；但如果你用scipy.ndimage.affine_transform或其他库，可能需要转置或调整顺序。最稳妥的方法是，用一个已知的简单变换（如纯平移[1,0,tx; 0,1,ty]）做测试，观察输出网格是否符合预期。

2.2 GridSample：拿着地图，精准“挖矿”

GridSample 是整个流程的执行者。它接收两个核心输入：原始图像/特征图input（shape:NxCxHxW）和 AffineGrid 生成的grid（shape:NxHxWx2）。它的任务，就是对grid中的每一个(x_norm, y_norm)坐标，在input上进行双线性（或最近邻）插值采样，并将结果填入输出张量的对应位置。

双线性插值是 GridSample 的默认模式，其原理非常直观：对于一个非整数的归一化坐标(x, y)，它找到周围四个最近的整数坐标点（即(floor(x), floor(y)),(ceil(x), floor(y)),(floor(x), ceil(y)),(ceil(x), ceil(y))），然后根据(x, y)到这四个点的距离，加权平均它们的像素值。权重由距离的线性函数决定，因此叫“双线性”。这个过程是完全可导的，梯度可以精确地反向传播回input和grid。这意味着，如果你的grid是由一个可学习的神经网络参数生成的，那么整个采样过程的梯度就能一路回传，驱动网络去学习如何生成最优的变形网格。

GridSample 的padding_mode参数决定了当grid中的坐标超出[-1, 1]范围时，如何处理。'zeros'（默认）会填 0；'border'会用边界像素值填充；'reflection'会像镜子一样反射坐标。这个选择在实际项目中影响巨大。例如，在训练一个图像修复模型时，如果变形导致大量区域超出原图边界，用'zeros'会引入大量黑色伪影，严重干扰损失函数；而'border'则能保持边缘的连贯性，让模型更专注于内容修复而非“学填黑”。

注意：GridSample 的输出 shape 与grid的 spatial dimensions（即HxW）完全一致，而不是与input的HxW一致。这是一个非常重要的设计。它意味着，你可以用一张64x64的网格，去采样一张256x256的大图，得到一张64x64的小图；反之，你也可以用一张256x256的网格，去采样一张64x64的小图，得到一张256x256的大图（此时会自动进行上采样）。这个特性是实现任意尺度变换的基础，也是它比传统resize函数更灵活的根本原因。

2.3 为什么是“兄弟”，而不是“父子”？

从 API 设计哲学上看，AffineGrid 和 GridSample 的分离，体现了 PyTorch “显式优于隐式”的核心思想。很多框架（如 TensorFlow 1.x 的tf.contrib.image.transform）提供了一个“一站式”函数，内部封装了从矩阵到采样的全过程。这种设计对初学者友好，但牺牲了灵活性。PyTorch 的方案则像一个精密的乐高系统：AffineGrid 是一块“变换定义砖”，GridSample 是一块“采样执行砖”。你可以自由组合：

• 用torch.eye(2, 3)生成一个恒等变换矩阵，喂给 AffineGrid，得到一个“直通”网格，再用 GridSample，就实现了无任何变形的复制。
• 用torch.tensor([[1, 0, 0.2], [0, 1, 0]])生成一个向右平移 20% 的矩阵，就能看到图像整体右移。
• 更进一步，你可以用一个小型 CNN，输入一张图像，输出一个2x3的theta，再经过 AffineGrid+GridSample，就构成了一个完整的 Spatial Transformer Network (STN) 模块。这个 CNN 的权重是可学习的，它学会了“看图说话”，自己决定该如何对图像进行预处理。

这种解耦带来的另一个巨大好处是调试友好。当你发现最终输出有问题时，你可以单独打印grid张量，用plt.imshow(grid[0, ..., 0])可视化 x 坐标通道，立刻就能判断出是变换矩阵写错了，还是 AffineGrid 的坐标系理解有误。如果所有逻辑都封装在一个黑盒函数里，这种定位将变得异常困难。所以，它们是“兄弟”，是平等协作的伙伴，共同构建了 PyTorch 空间操作的坚实底座。

3. 实操全流程：从零开始，亲手搭建一个可学习的图像变形器

3.1 环境准备与数据加载：一个最小但完整的起点

我们不依赖任何外部数据集，直接用torchvision生成一个合成的、带明显几何特征的测试图像。这样做的好处是，结果完全可控，便于验证每一步的正确性。首先，确保你的环境已安装 PyTorch（建议 2.0+）和 torchvision。

pip install torch torchvision matplotlib numpy

接下来，我们创建一个 128x128 的“棋盘格”图像，它由 8x8 个 16x16 的黑白方块组成。这种高对比度、强几何结构的图像，是检验空间变换效果的黄金标准，任何微小的变形、插值瑕疵都会被立刻暴露。

import torch import torch.nn as nn import torch.nn.functional as F import numpy as np import matplotlib.pyplot as plt def create_checkerboard(h=128, w=128, block_size=16): """生成一个 h x w 的棋盘格图像，返回 [1, 1, h, w] 的 float32 张量""" board = np.zeros((h, w), dtype=np.float32) for i in range(0, h, block_size): for j in range(0, w, block_size): if (i // block_size + j // block_size) % 2 == 0: board[i:i+block_size, j:j+block_size] = 1.0 return torch.from_numpy(board).unsqueeze(0).unsqueeze(0) # [1, 1, h, w] # 创建原始图像 original_img = create_checkerboard() print(f"原始图像 shape: {original_img.shape}") # torch.Size([1, 1, 128, 128])

现在，我们有了一个干净的、可复现的输入。下一步，是定义我们的“可学习变形器”的核心——一个能输出theta矩阵的神经网络。它不需要复杂，一个两层全连接网络足矣，但必须包含一个tanh激活，将输出限制在[-1, 1]附近，防止初始权重过大导致grid坐标严重越界。

class SimpleSTN(nn.Module): def __init__(self, input_channels=1): super().__init__() # 输入是 [1, 1, 128, 128]，我们先用一个简单的卷积提取全局特征 self.conv = nn.Sequential( nn.Conv2d(input_channels, 8, kernel_size=3, padding=1), nn.ReLU(), nn.AdaptiveAvgPool2d((1, 1)) # 全局平均池化，得到 [1, 8, 1, 1] ) # 将 8 维特征映射到 6 维的 theta (2x3) self.fc = nn.Sequential( nn.Linear(8, 32), nn.ReLU(), nn.Linear(32, 6) # 输出 6 个值，用于 reshape 成 2x3 ) def forward(self, x): # x: [N, C, H, W] feat = self.conv(x).view(x.size(0), -1) # [N, 8] theta_flat = self.fc(feat) # [N, 6] # reshape 成 [N, 2, 3] theta = theta_flat.view(-1, 2, 3) # 初始化为恒等变换，然后让网络学习微调 identity = torch.tensor([[[1, 0, 0], [0, 1, 0]]], dtype=torch.float32, device=x.device) theta = identity + 0.1 * theta # 小幅扰动，避免训练初期崩溃 return theta # 实例化模型 stn = SimpleSTN() # 生成一个 theta 矩阵用于测试 theta_test = stn(original_img) print(f"生成的 theta shape: {theta_test.shape}") # torch.Size([1, 2, 3]) print(f"theta 值:\n{theta_test[0]}")

这段代码的关键在于theta的初始化策略。我们没有让它从随机值开始，而是以恒等变换[1,0,0; 0,1,0]为基础，加上一个0.1 * theta的小扰动。这是非常重要的工程技巧。如果theta初始值过大，AffineGrid 生成的grid会大面积越界，GridSample 在'zeros'模式下会输出大片黑色，导致梯度消失，模型根本无法学习。这个0.1的系数，就是我们给模型的“温柔提示”，告诉它：“先学会微调，再学大变”。

3.2 核心变换链：AffineGrid + GridSample 的完整调用

现在，我们把前面定义的theta送入 AffineGrid，再将生成的grid与原始图像一起送入 GridSample。这是整个流程的“心脏”。

# Step 1: 使用 AffineGrid 生成采样网格 # grid 的 shape 是 [N, H_out, W_out, 2] # 这里我们让输出尺寸与输入相同，即 128x128 grid = F.affine_grid(theta_test, size=original_img.size(), align_corners=True) print(f"AffineGrid 生成的 grid shape: {grid.shape}") # torch.Size([1, 128, 128, 2]) print(f"grid x 坐标范围: [{grid[0, ..., 0].min():.3f}, {grid[0, ..., 0].max():.3f}]") print(f"grid y 坐标范围: [{grid[0, ..., 1].min():.3f}, {grid[0, ..., 1].max():.3f}]") # Step 2: 使用 GridSample 执行采样 # output 的 shape 与 grid 的 spatial dims 一致，即 [1, 1, 128, 128] output_img = F.grid_sample(original_img, grid, mode='bilinear', padding_mode='zeros', align_corners=True) print(f"GridSample 输出 shape: {output_img.shape}")

这里有两个至关重要的参数：align_corners和padding_mode。align_corners=True是 PyTorch 1.3 之后的默认行为，它保证了归一化坐标-1和1分别精确对应源图像的左上角和右下角像素中心。如果你设为False，坐标-1和1会对应到像素的边界，这会导致轻微的缩放偏差，在高精度任务中必须统一。padding_mode='zeros'是最安全的默认选项，但在实际项目中，你可能需要根据任务切换。例如，在训练一个图像超分辨率模型时，'border'通常能产生更自然的边缘。

实操心得：在调试阶段，务必打印grid的坐标范围。如果x或y的最大值远大于1或最小值远小于-1，说明你的theta矩阵数值过大，需要检查网络初始化或学习率。一个健康的grid，其坐标范围应该紧密围绕[-1, 1]，比如[-1.05, 1.05]是可以接受的，但[-3.2, 4.7]就是灾难性的。

3.3 可视化与验证：用眼睛“读懂”坐标变换

光看 print 输出是不够的，我们必须用可视化来建立直觉。我们将原始图像、grid的 x/y 分量、以及最终输出，四宫格并排显示。

def visualize_transformation(original, grid, output): fig, axes = plt.subplots(2, 2, figsize=(10, 10)) # 原始图像 axes[0, 0].imshow(original[0, 0].cpu().numpy(), cmap='gray') axes[0, 0].set_title('Original') axes[0, 0].axis('off') # grid x 分量 axes[0, 1].imshow(grid[0, ..., 0].cpu().numpy(), cmap='RdBu_r', vmin=-1, vmax=1) axes[0, 1].set_title('Grid X (Normalized)') axes[0, 1].axis('off') # grid y 分量 axes[1, 0].imshow(grid[0, ..., 1].cpu().numpy(), cmap='RdBu_r', vmin=-1, vmax=1) axes[1, 0].set_title('Grid Y (Normalized)') axes[1, 0].axis('off') # 输出图像 axes[1, 1].imshow(output[0, 0].cpu().numpy(), cmap='gray') axes[1, 1].set_title('Output (Transformed)') axes[1, 1].axis('off') plt.tight_layout() plt.show() # 执行可视化 visualize_transformation(original_img, grid, output_img)

这个可视化是调试的利器。观察Grid X和Grid Y图：它们应该是平滑变化的彩色图，颜色从蓝色（-1）渐变到红色（1）。如果其中一张图是纯蓝色或纯红色，说明该方向的坐标没有变化，即没有发生平移或缩放。如果图中出现了尖锐的、不连续的色块，那很可能theta矩阵的某个元素是nan或inf，需要检查梯度爆炸。而最终的Output图，应该是一个清晰、无撕裂的棋盘格。如果出现了模糊、重影或黑色条纹，问题一定出在grid的生成或GridSample的参数上。

3.4 构建端到端训练循环：让网络学会“看图变形”

现在，我们把所有部件组装成一个可训练的模型。我们将定义一个简单的损失函数：让变形后的图像，尽可能接近一个“理想”的目标图像。这个目标图像，我们可以人为制造——比如，让棋盘格向右平移 10 个像素。

# 定义一个“理想”的目标变换：向右平移 10 像素 # 在 128x128 图像中，10 像素占 10/127 ≈ 0.0787 的归一化距离 target_theta = torch.tensor([[[1, 0, 0.0787], [0, 1, 0]]], dtype=torch.float32) # 生成目标图像（用于监督） target_grid = F.affine_grid(target_theta, size=original_img.size(), align_corners=True) target_output = F.grid_sample(original_img, target_grid, mode='bilinear', padding_mode='zeros', align_corners=True) # 定义模型和优化器 model = SimpleSTN() optimizer = torch.optim.Adam(model.parameters(), lr=1e-3) # 训练循环 for epoch in range(100): optimizer.zero_grad() # 前向传播 theta_pred = model(original_img) grid_pred = F.affine_grid(theta_pred, size=original_img.size(), align_corners=True) output_pred = F.grid_sample(original_img, grid_pred, mode='bilinear', padding_mode='zeros', align_corners=True) # 计算损失：L1 损失 loss = F.l1_loss(output_pred, target_output) # 反向传播 loss.backward() optimizer.step() if epoch % 20 == 0: print(f"Epoch {epoch}, Loss: {loss.item():.6f}") # 训练完成后，再次可视化 theta_final = model(original_img) grid_final = F.affine_grid(theta_final, size=original_img.size(), align_corners=True) output_final = F.grid_sample(original_img, grid_final, mode='bilinear', padding_mode='zeros', align_corners=True) visualize_transformation(original_img, grid_final, output_final)

这个训练循环虽然简单，但它完整展示了 AffineGrid 和 GridSample 如何融入一个可学习的 pipeline。关键点在于，loss.backward()会自动计算loss对theta_pred的梯度，而theta_pred的梯度又会通过F.affine_grid和F.grid_sample的可导性，反向传播到model的所有参数上。这就是“端到端可微”的魔力。运行这个代码，你会发现，随着 epoch 增加，output_final会越来越接近target_output，theta_final的t_x分量也会逐渐收敛到0.0787。这证明了整个链条是健壮且可学习的。

4. 高阶应用与避坑指南：那些文档里不会写的实战经验

4.1 场景一：批量处理不同尺寸的图像——动态 grid 尺寸的陷阱

在真实项目中，你的输入图像 rarely 是固定尺寸的。你可能会遇到一个 batch 里，有的图像是128x128，有的是256x256。这时，F.affine_grid的size参数就不能写死。一个常见的错误是，试图用original_img.size()直接传入，但这会返回[N, C, H, W]，而affine_grid需要的是[N, C, H_out, W_out]。如果你的 batch 内尺寸不一致，original_img.size()会报错。

正确做法是：为每个样本单独生成grid。你可以用一个 for 循环，但这会损失性能。更优雅的方式是，利用 PyTorch 的广播机制，预先计算好每个样本的目标尺寸，然后用torch.stack组合。

# 假设我们有一个 batch，尺寸各不相同 batch_imgs = [ create_checkerboard(128, 128), # [1, 1, 128, 128] create_checkerboard(256, 192), # [1, 1, 256, 192] create_checkerboard(64, 64), # [1, 1, 64, 64] ] # 将它们拼成一个 batch: [3, 1, H_i, W_i]，但这是不合法的，因为 H_i/W_i 不同 # 所以，我们用 list 存储，然后分别处理 grids = [] outputs = [] thetas = [] for i, img in enumerate(batch_imgs): # 为每个 img 单独生成 theta theta_i = stn(img) # theta_i shape: [1, 2, 3] thetas.append(theta_i) # 动态指定 size size_i = img.size() # [1, 1, H, W] grid_i = F.affine_grid(theta_i, size=size_i, align_corners=True) # [1, H, W, 2] grids.append(grid_i) output_i = F.grid_sample(img, grid_i, mode='bilinear', padding_mode='zeros', align_corners=True) outputs.append(output_i) # 如果你需要一个统一的 tensor，可以考虑 padding 到相同尺寸，但这会引入额外的空白 # 或者，接受 list of tensors 的形式，这在 DataLoader 的 collate_fn 中很常见

避坑心得：永远不要假设 batch 内所有样本尺寸相同。在写通用代码时，把size参数当作一个需要动态计算的变量，而不是一个静态常量。一个实用的技巧是，在DataLoader的collate_fn中，为每个 batch 返回一个sizes列表，与batch_tensor并行传递。

4.2 场景二：处理多通道特征图——RGB 与特征图的差异

F.grid_sample的input参数，可以是任意NxCxHxW的张量。这包括 RGB 图像（C=3），也包括 CNN 中间层的特征图（C=64, 128, ...）。但这里有一个微妙的差异：对于 RGB 图像，我们通常关心视觉保真度，所以mode='bilinear'是首选；而对于高维特征图，mode='bilinear'可能会引入不必要的平滑，破坏特征的判别性。此时，mode='nearest'反而更合适，因为它只是做整数索引，没有插值，能最大程度保留原始特征的“锐度”。

# 假设我们有一个 4D 特征图，来自 ResNet 的某一层 feature_map = torch.randn(1, 64, 32, 32) # [1, 64, 32, 32] # 生成一个缩放 1.5 倍的 theta scale = 1.5 theta_scale = torch.tensor([[[scale, 0, 0], [0, scale, 0]]], dtype=torch.float32) # 为特征图生成 grid，注意 size 是特征图的尺寸 grid_feat = F.affine_grid(theta_scale, size=feature_map.size(), align_corners=True) # 使用 nearest 模式采样，避免特征模糊 output_feat = F.grid_sample(feature_map, grid_feat, mode='nearest', padding_mode='zeros', align_corners=True) print(f"特征图变换前 shape: {feature_map.shape}") # [1, 64, 32, 32] print(f"特征图变换后 shape: {output_feat.shape}") # [1, 64, 32, 32] (因为 grid 是 32x32)

实操心得：在构建 STN 模块时，如果你的 STN 是插在 CNN 主干的中间，用来对特征图进行空间校准，那么mode='nearest'通常是更好的默认选择。只有当你明确需要亚像素级别的精细变形（如光流估计），才使用bilinear。

4.3 场景三：调试与诊断——当输出全是黑/白/模糊时怎么办？

这是新手最常遇到的“玄学”问题。下面是一份基于我多年 debug 经验的速查表：

一个终极的 debug 技巧是：绕过 AffineGrid，手动构造一个已知的、简单的grid。例如，构造一个让图像向右平移 1 个像素的grid：

# 手动构造一个平移 grid (向右 1 像素) h, w = 128, 128 yy, xx = torch.meshgrid(torch.linspace(-1, 1, h), torch.linspace(-1, 1, w), indexing='ij') # xx 是 x 坐标，向右平移 1 像素，相当于在归一化坐标中增加 2/(w-1) delta_x = 2.0 / (w - 1) grid_manual = torch.stack([xx + delta_x, yy], dim=-1).unsqueeze(0) # [1, h, w, 2] output_manual = F.grid_sample(original_img, grid_manual, mode='bilinear', padding_mode='zeros', align_corners=True)

如果这个手动grid能得到正确的平移效果，那就 100% 证明F.grid_sample本身工作正常，问题一定出在F.affine_grid的theta输入上。这个“隔离法”能帮你瞬间锁定问题域，节省大量时间。

4.4 性能优化：GPU 上的极致加速

在大规模训练中，F.affine_grid和F.grid_sample的性能至关重要。它们的 CUDA kernel 已经高度优化，但仍有几个关键点可以提升：

1. Batch Size 是王道：这两个函数的 GPU 利用率与 batch size 正相关。单个样本（N=1）的吞吐量远低于N=32。如果你的应用允许，尽量增大 batch size。
2. 内存连续性：确保input和grid张量都是contiguous()的。非连续张量会触发一次隐式的内存拷贝，带来显著开销。在forward函数末尾，可以加一句grid = grid.contiguous()。
3. 避免重复计算：如果你的theta在一个 batch 内是相同的（例如，对整个 batch 做统一的随机增强），那么F.affine_grid只需调用一次，然后用expand复制到整个 batch，而不是对每个样本都调用。

# 错误：对每个样本都调用 affine_grid grids_bad = [] for i in range(N): grid_i = F.affine_grid(theta_single, size=input_batch[i].size()) grids_bad.append(grid_i)