PyTorch中针对图像块的可微仿射变换实现方案咨询

原代码问题说明 #

你遇到的原地操作报错核心原因是直接对输入图像张量做img[im_id, :, x1:x2, y1:y2] = ...的原地赋值，这类inplace操作会打断PyTorch的计算图回溯，导致反向传播失败。for循环本身不会影响可微性，只要替换成非原地的张量操作即可，也可以搭配affine_grid+grid_sample（即STN的核心接口）实现完全可微的逻辑。

可微实现思路 #

全程避免原地修改输入张量，通过张量克隆+index_put非原地赋值的方式保留完整计算图，逻辑和你原有的循环实现完全对齐：

1. 先克隆原图像张量作为输出基底，避免修改原始输入
2. 循环处理每个物体的边界框和仿射参数，用grid_sample生成变换后的图像块
3. 用index_put接口将变换后的块写入输出张量的对应位置，该接口返回新张量而非修改原张量，全程可微

代码实现 #

import torch
import torch.nn.functional as F

# 输入参数说明：
# imgs: 批量图像，形状[B, C, H, W]，对应你的输入为[2,3,64,64]
# boxes_pred: 物体边界框，形状[N,4]，对应你的输入为[10,4]，格式为(x1,y1,x2,y2)像素坐标
# theta_mean: 仿射变换参数，形状[N,6]，每个对应一个物体的2x3仿射矩阵展平结果
# obj_to_img: 物体归属向量，形状[N]，每个值为对应图像的batch维度索引

B, C, H, W = imgs.shape
N = boxes_pred.shape[0]
device = imgs.device

# 克隆原图作为输出基底，避免原地修改原始输入
output = imgs.clone()

for obj_num in range(N):
    im_id = obj_to_img[obj_num]
    # 边界框坐标转整数索引，同时裁剪到图像有效范围内避免越界
    x1 = max(0, int(boxes_pred[obj_num, 0]))
    y1 = max(0, int(boxes_pred[obj_num, 1]))
    x2 = min(H, int(boxes_pred[obj_num, 2]))
    y2 = min(W, int(boxes_pred[obj_num, 3]))
    h_patch = x2 - x1
    w_patch = y2 - y1
    # 过滤无效边界框
    if h_patch <= 0 or w_patch <= 0:
        continue
    
    # 取出对应原始图像块
    im_patch = imgs[im_id:im_id+1, :, x1:x2, y1:y2]
    # 仿射参数调整为affine_grid要求的[1,2,3]形状
    theta = theta_mean[obj_num].view(1, 2, 3)
    
    # 生成仿射采样网格，align_corners参数和你原有STN实现保持一致
    grid = F.affine_grid(theta, im_patch.size(), align_corners=False)
    # 采样得到变换后的图像块
    transformed_patch = F.grid_sample(im_patch, grid, align_corners=False, padding_mode='zeros')
    
    # 非原地将变换后的块写入输出张量对应位置
    output = output.index_put(
        indices=(im_id, slice(None), slice(x1, x2), slice(y1, y2)),
        values=transformed_patch[0],
        accumulate=False
    )

注意事项 #

• 如果所有边界框的尺寸一致，可以进一步将循环逻辑改为批量grid_sample操作，提升运行效率
• align_corners参数需要和你原有STN实现的参数对齐，避免变换结果出现偏移
• 若存在多物体边界框重叠的情况，默认后处理的物体块会覆盖先处理的块，和你原有逻辑一致；如果需要处理重叠，可以额外加掩码加权平均的逻辑

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