求PyTorch中对应TensorFlow tfa.image.transform的等价函数

1. 仿射变换场景（对应基础 affine 矩阵） #

如果你的变换是仿射变换（平移、旋转、缩放、剪切的组合），可以直接用 torchvision.transforms.functional.affine，它会自动处理变换矩阵，用法如下：

import torch
import torchvision.transforms.functional as F

# 示例：旋转30度，平移(10,10)，缩放0.8，无剪切
img_tensor = torch.randn(1, 3, 256, 256)  # 张量形状: [batch, channels, H, W]
transformed_img = F.affine(
    img_tensor,
    angle=30,
    translate=(10, 10),
    scale=0.8,
    shear=(0, 0)
)

2. 自定义任意变换矩阵（完全匹配 tfa.image.transform 的通用场景） #

tfa.image.transform 支持传入任意3x3齐次变换矩阵，PyTorch中可以通过 torch.nn.functional.grid_sample 配合手动生成变换网格实现，步骤如下：

import torch
import torch.nn.functional as F

def transform_image(img, transform_matrix):
    """
    实现与tfa.image.transform一致的自定义矩阵变换逻辑
    :param img: 输入图像张量，形状[batch, channels, H, W]
    :param transform_matrix: 3x3齐次变换矩阵，支持[batch, 3, 3]或[3,3]格式
    :return: 变换后的图像张量
    """
    batch_size, channels, h, w = img.shape
    # 生成标准化坐标网格，范围[-1,1]（PyTorch grid_sample要求的输入格式）
    grid_y, grid_x = torch.meshgrid(torch.linspace(-1, 1, h), torch.linspace(-1, 1, w))
    grid = torch.stack([grid_x, grid_y], dim=-1)  # 形状(H,W,2)
    grid = grid.unsqueeze(0).repeat(batch_size, 1, 1, 1)  # 扩展到batch维度
    
    # 转换为齐次坐标
    grid_homogeneous = torch.cat([grid, torch.ones_like(grid[..., :1])], dim=-1)
    
    # 应用变换矩阵
    if transform_matrix.dim() == 2:
        transform_matrix = transform_matrix.unsqueeze(0).repeat(batch_size, 1, 1)
    transformed_grid_homogeneous = torch.matmul(grid_homogeneous, transform_matrix.transpose(1,2))
    
    # 还原为2D坐标并截断到[-1,1]，避免超出图像区域的异常
    transformed_grid = transformed_grid_homogeneous[..., :2] / transformed_grid_homogeneous[..., 2:3]
    transformed_grid = torch.clamp(transformed_grid, -1, 1)
    
    # 执行图像变换，mode和padding_mode可按需调整，对应tfa的插值和填充方式
    return F.grid_sample(img, transformed_grid, mode='bilinear', padding_mode='border', align_corners=True)

# 示例：构造平移+旋转的变换矩阵
theta = torch.tensor([
    [0.866, -0.5, 0.2],  # cos30°、-sin30°、x方向平移量
    [0.5, 0.866, 0.1],   # sin30°、cos30°、y方向平移量
    [0, 0, 1]
])
img_batch = torch.randn(2, 3, 256, 256)  # 2张图像的batch
transformed_batch = transform_image(img_batch, theta)

关键注意点 #

• 设置 grid_sample 的 align_corners=True，能让变换对齐方式更接近TensorFlow的逻辑
• 单张图像处理时，可先通过 img.unsqueeze(0) 添加batch维度，变换后用 result.squeeze(0) 还原

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