PyTorch中transform函数无法被autograd求导的问题咨询

使用torchvision的图像变换后，调用normalized.backward()时，resized和transformed张量均无梯度产生，即使单独对每个变换求梯度也失败。相关代码如下：

from torchvision import transforms
from torchvision.transforms import ToTensor

resize = transforms.Resize(size=224, interpolation=transforms.InterpolationMode.BICUBIC, max_size=None, antialias=None)
crop = transforms.CenterCrop(size=(224, 224))
normalize = transforms.Normalize(mean=(0.48145466, 0.4578275, 0.40821073), std=(0.26862954, 0.26130258, 0.27577711))

img = torch.Tensor(images[30])
img.requires_grad = True

rgb = torch.dsplit(torch.Tensor(img),3)
transformed = torch.stack(rgb).reshape(3,100,100)

resized = resize.forward(transformed)
normalized = normalize.forward(resized)

image_features = clip_model.encode_image(normalized.unsqueeze(0).to(device))
text_features = clip_model.encode_text(text_inputs)
similarity = (100.0 * image_features @ text_features.T).softmax(dim=-1)

1. 修复计算图断开的核心问题 #

代码中rgb = torch.dsplit(torch.Tensor(img),3)这一步重新创建了默认requires_grad=False的新张量，直接断开了原img的梯度追踪链路。应直接对已开启梯度的原张量操作：

# 直接创建带梯度的初始张量，无需后续单独设置
img = torch.tensor(images[30], dtype=torch.float32, requires_grad=True)
rgb = torch.dsplit(img, 3)  # 基于原张量拆分，保留梯度追踪
transformed = torch.stack(rgb).reshape(3, 100, 100)

2. 确保Resize变换的可微分性 #

PyTorch的Resize在BICUBIC插值模式下支持微分，但需注意：

• 使用PyTorch 1.10及以上版本
• 显式设置antialias=True（不同版本默认值有差异，显式设置可保证梯度计算稳定性）：

resize = transforms.Resize(
    size=224,
    interpolation=transforms.InterpolationMode.BICUBIC,
    antialias=True
)

3. 规范变换调用方式 #

无需使用.forward()方法，直接调用变换对象即可，避免潜在的计算图异常：

resized = resize(transformed)
normalized = normalize(resized)

4. 保证设备一致性 #

若模型在GPU运行，初始张量需同步移至GPU，避免跨设备拷贝导致梯度丢失：

img = torch.tensor(images[30], dtype=torch.float32, requires_grad=True).to(device)

修正后的完整代码示例 #

from torchvision import transforms
import torch

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

resize = transforms.Resize(
    size=224,
    interpolation=transforms.InterpolationMode.BICUBIC,
    antialias=True
)
normalize = transforms.Normalize(
    mean=(0.48145466, 0.4578275, 0.40821073),
    std=(0.26862954, 0.26130258, 0.27577711)
)

# 创建带梯度的初始张量并移至目标设备
img = torch.tensor(images[30], dtype=torch.float32, requires_grad=True).to(device)
rgb = torch.dsplit(img, 3)
transformed = torch.stack(rgb).reshape(3, 100, 100)

# 执行变换
resized = resize(transformed)
normalized = normalize(resized)

# 后续模型计算
image_features = clip_model.encode_image(normalized.unsqueeze(0))
text_features = clip_model.encode_text(text_inputs)
similarity = (100.0 * image_features @ text_features.T).softmax(dim=-1)

# 针对具体损失值反向传播（直接对normalized调用backward不符合常规梯度计算逻辑）
similarity[0, 0].backward()
print(img.grad)  # 此时可正常获取初始张量的梯度

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