如何为Torch张量实现Image.AFFINE的等效功能？

原代码问题说明 #

你原来的代码中，frame.transform()是PIL Image的专属方法，但torch.load()加载出的是Torch张量，无法直接调用该方法。原逻辑是通过水平平移offset像素实现仿射变换（AFFINE参数(1, 0, offset, 0, 1, 0)对应仅x轴平移，y轴无变化），下面提供两种PyTorch原生的替代实现方案。

方案1：用torchvision快速实现（推荐） #

torchvision.transforms.functional.affine()直接支持Torch张量输入，完全匹配原逻辑：

import torchvision.transforms.functional as F
import random

# 原代码替换部分
frame = torch.load(file_path + '/' + f)            
offset += random.randrange(-v, v)
offset = min(offset, 3 * v)
offset = max(offset, -3 * v)

# 替换PIL的AFFINE变换为PyTorch版本
# 参数对应：角度0（无旋转）、平移(x=offset, y=0)、缩放1、剪切0
frame = F.affine(frame, angle=0, translate=(offset, 0), scale=1.0, shear=0)

if self.transform is not None:
    frame = self.transform[0](frame)
frames.append(frame)

方案2：手动构建仿射矩阵（底层实现） #

如果不想依赖torchvision，可通过affine_grid和grid_sample手动实现：

import torch.nn.functional as F
import random

# 原代码替换部分
frame = torch.load(file_path + '/' + f)            
offset += random.randrange(-v, v)
offset = min(offset, 3 * v)
offset = max(offset, -3 * v)

# 处理张量维度：单张图(C,H,W)转成批量格式(1,C,H,W)
is_single = len(frame.shape) == 3
if is_single:
    frame = frame.unsqueeze(0)

# 构建仿射矩阵：PyTorch要求平移量归一化到[-1,1]范围
batch_size, _, _, img_width = frame.shape
tx = offset / img_width * 2  # 像素偏移转归一化坐标
affine_mat = torch.tensor([[1.0, 0.0, tx], [0.0, 1.0, 0.0]], device=frame.device)
affine_mat = affine_mat.unsqueeze(0).repeat(batch_size, 1, 1)

# 生成采样网格并执行变换
grid = F.affine_grid(affine_mat, frame.size(), align_corners=False)
frame = F.grid_sample(frame, grid, align_corners=False)

# 还原单张图维度
if is_single:
    frame = frame.squeeze(0)

if self.transform is not None:
    frame = self.transform[0](frame)
frames.append(frame)

注意事项 #

• 张量维度要求：确保输入张量为(C, H, W)或(B, C, H, W)格式，如果是(H, W, C)需先转置：frame = frame.permute(2, 0, 1)
• 数值范围：若张量是uint8类型（0-255），建议先转成float32再做变换：frame = frame.float() / 255.0，变换后可按需转回原类型

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