PyTorch分割任务中如何为图像与掩膜实现同步的transforms.Compose配对变换

核心逻辑是把所有几何类空间变换和非空间的像素值变换拆分，通过固定随机生成状态的方式保证图像和掩膜的空间变换参数完全对齐，同时给两类样本配置各自专属的非空间处理逻辑。

核心问题根因 #

原来的实现方式里，哪怕给图像和掩膜配的随机变换参数完全一致，两条独立流水线在生成随机决策的时候（比如随机掷硬币决定要不要水平翻转、随机抽一个旋转角度、随机框一个裁剪区域）是独立从随机数生成器取数的，没有任何参数共享逻辑，最后得到的图像和掩膜自然会出现空间错位。而高斯噪声是直接修改像素值的操作，根本不涉及空间位置变化，本身就不需要和掩膜做同步。

分步实现步骤 #

1. 拆分两类变换逻辑 #

把所有变换按是否涉及空间几何操作拆分：

• 空间几何类（必须完全同步）：Resize(512)、p=0.5的RandomHorizontalFlip、p=0.5的RandomVerticalFlip、RandomRotation(90)、scale为(0.3,1.0)的RandomResizedCrop(320)
• 像素非空间类（各自独立，不需要同步）：
  • 图像专属：自定义AddGaussianNoise(0.,1.)、ToTensor、适配ImageNet均值/标准差的Normalize
  • 掩膜专属：仅需将掩膜转为LongTensor格式的张量，不需要加噪声、不需要做均值标准差归一化

常见踩坑提醒：掩膜做几何变换时，所有插值操作必须使用最近邻模式，不能用默认的双线性插值，否则会将整数类别的掩膜值插值成小数，导致标签索引错乱。
因此需要给图像和掩膜分别配置同随机参数、不同插值模式的空间流水线：

from torchvision import transforms
import random
import numpy as np
import torch

# 图像侧空间流水线（双线性插值，适合RGB图像）
image_spatial_tf = transforms.Compose([
    transforms.Resize(512),
    transforms.RandomHorizontalFlip(p=0.5),
    transforms.RandomVerticalFlip(p=0.5),
    transforms.RandomRotation(90),
    transforms.RandomResizedCrop(320, scale=(0.3, 1.0))
])

# 掩膜侧空间流水线（最近邻插值，参数和图像侧完全对齐）
mask_spatial_tf = transforms.Compose([
    transforms.Resize(512, interpolation=transforms.InterpolationMode.NEAREST),
    transforms.RandomHorizontalFlip(p=0.5),
    transforms.RandomVerticalFlip(p=0.5),
    transforms.RandomRotation(90, interpolation=transforms.InterpolationMode.NEAREST),
    transforms.RandomResizedCrop(320, scale=(0.3, 1.0), interpolation=transforms.InterpolationMode.NEAREST)
])

# 图像专属非空间变换
image_only_tf = transforms.Compose([
    AddGaussianNoise(0., 1.), # 原项目中实现的高斯噪声增强类
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

# 掩膜专属非空间变换
mask_only_tf = transforms.Compose([
    transforms.ToTensor(),
    transforms.Lambda(lambda x: x.squeeze(0).long()) # 保证掩膜为LongTensor格式
])

2. 用随机状态保存/恢复保证参数完全同步 #

很多人会踩的误区是：直接把同一个空间Compose先后传给图像和掩膜做处理。实际上几何变换的随机参数是在每次调用时动态生成的，处理图像的过程会连续消耗多次随机数，等轮到处理掩膜时，随机数生成器的状态已经偏移，两次生成的参数根本对不上。
正确的处理流程是：在对图像做空间变换前，先保存全局所有随机源（Python内置random、numpy、torch）的当前状态，处理完图像后把随机状态恢复到保存的节点，再对掩膜执行空间变换，此时两次生成的所有随机参数（是否翻转、旋转角度、裁剪框坐标）完全一致。
实现代码：

def process_paired_sample(raw_image, raw_mask):
    # 保存全链路随机状态
    rng_py = random.getstate()
    rng_np = np.random.get_state()
    rng_torch = torch.random.get_rng_state()

    # 对图像执行空间变换
    image_spatial = image_spatial_tf(raw_image)

    # 恢复随机状态到之前保存的节点
    random.setstate(rng_py)
    np.random.set_state(rng_np)
    torch.random.set_rng_state(rng_torch)

    # 对掩膜执行同参数的空间变换
    mask_spatial = mask_spatial_tf(raw_mask)

    # 分别执行各自的非空间专属变换
    final_image = image_only_tf(image_spatial)
    final_mask = mask_only_tf(mask_spatial)

    return final_image, final_mask

3. 训练流程接入方式 #

把原来Dataset的__getitem__方法里，图像、掩膜单独调用transform的逻辑，替换成上面的process_paired_sample方法即可：

• 所有空间类操作在随机状态绑定的流程下完成，100%保证图像和掩膜的空间位置对齐
• 高斯噪声操作在空间变换完成后单独对图像执行，不会干扰掩膜的变换逻辑，也不会打乱随机状态导致对齐失败
• 不需要重写所有transform的内部逻辑，改动量极小。

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