如何实现支持PyTorch梯度流的图像缩放插值函数？

核心问题分析 #

你遇到的问题根源有两个：

• F.interpolate若使用基于scale_factor计算的整数尺寸（size参数），会因整数转换阻断梯度流；直接传入scale_factor张量时，部分PyTorch版本对动态缩放因子的梯度支持有限。
• 自定义双线性插值中，int()、.long()这类离散类型转换和整数索引操作会破坏梯度传递链，导致scale_factor的梯度无法反向传播。

解决方案1：基于grid_sample的可导缩放（推荐） #

torch.nn.functional.grid_sample原生支持梯度传递，只要生成的采样网格是基于scale_factor的可导计算结果。以下是完整实现：

import torch
import torch.nn.functional as F

def differentiable_interpolate(img, scale_factor):
    """
    支持梯度流的图像缩放（双线性插值）
    Args:
        img: 输入张量，形状[N, C, H, W]
        scale_factor: 缩放因子张量（带梯度），形状[]或[N]
    Returns:
        缩放后的图像张量，形状[N, C, H*scale_factor, W*scale_factor]
    """
    n, c, h, w = img.size()
    device = img.device

    # 计算目标尺寸（整数），保留scale_factor的浮点关联
    new_h = torch.round(h * scale_factor).long() if isinstance(scale_factor, torch.Tensor) else int(h * scale_factor)
    new_w = torch.round(w * scale_factor).long() if isinstance(scale_factor, torch.Tensor) else int(w * scale_factor)

    # 生成采样网格（归一化到[-1,1]，符合grid_sample要求）
    x = torch.linspace(0, w-1, new_w, device=device) / (w-1) * 2 - 1
    y = torch.linspace(0, h-1, new_h, device=device) / (h-1) * 2 - 1
    grid_y, grid_x = torch.meshgrid(y, x, indexing='ij')
    grid = torch.stack([grid_x, grid_y], dim=-1)  # 形状[new_h, new_w, 2]

    # 适配batch维度的缩放因子
    if scale_factor.dim() > 0:
        grid = grid.unsqueeze(0).repeat(n, 1, 1, 1)

    # 执行双线性插值，align_corners=True保证与F.interpolate行为一致
    interpolated = F.grid_sample(img, grid, mode='bilinear', align_corners=True)
    return interpolated

可行性说明 #

采样网格grid基于scale_factor计算的连续浮点值生成，所有操作均为可导运算，梯度可从插值结果反向传播到scale_factor；grid_sample是PyTorch原生优化实现，稳定性和效率远高于自定义插值。

解决方案2：修复自定义双线性插值 #

若坚持使用自定义实现，需修改类型转换和索引方式，避免阻断梯度：

def differentiable_bilinear_interpolate(img, scale_factor):
    n, c, h, w = img.size()
    device = img.device

    # 计算目标尺寸（整数）
    new_h = int(torch.round(h * scale_factor).item()) if isinstance(scale_factor, torch.Tensor) else int(h * scale_factor)
    new_w = int(torch.round(w * scale_factor).item()) if isinstance(scale_factor, torch.Tensor) else int(w * scale_factor)

    # 生成输入坐标网格（浮点，可导）
    h_coords = torch.linspace(0, h-1, new_h, device=device)
    w_coords = torch.linspace(0, w-1, new_w, device=device)
    grid_h, grid_w = torch.meshgrid(h_coords, w_coords, indexing='ij')

    # 计算四个邻域点的坐标（浮点，不转整数）
    h_floor = torch.floor(grid_h)
    h_ceil = h_floor + 1
    h_ceil = torch.clamp(h_ceil, max=h-1)
    w_floor = torch.floor(grid_w)
    w_ceil = w_floor + 1
    w_ceil = torch.clamp(w_ceil, max=w-1)

    # 计算插值权重（可导）
    h_weight = grid_h - h_floor
    w_weight = grid_w - w_floor

    # 用gather提取邻域点（支持梯度传递）
    img_flat = img.view(n, c, h*w)
    tl_idx = (h_floor.long() * w + w_floor.long()).view(1,1,new_h*new_w).repeat(n,c,1)
    tr_idx = (h_floor.long() * w + w_ceil.long()).view(1,1,new_h*new_w).repeat(n,c,1)
    bl_idx = (h_ceil.long() * w + w_floor.long()).view(1,1,new_h*new_w).repeat(n,c,1)
    br_idx = (h_ceil.long() * w + w_ceil.long()).view(1,1,new_h*new_w).repeat(n,c,1)

    tl = torch.gather(img_flat, 2, tl_idx).view(n,c,new_h,new_w)
    tr = torch.gather(img_flat, 2, tr_idx).view(n,c,new_h,new_w)
    bl = torch.gather(img_flat, 2, bl_idx).view(n,c,new_h,new_w)
    br = torch.gather(img_flat, 2, br_idx).view(n,c,new_h,new_w)

    # 双线性加权计算
    top = tl * (1 - w_weight) + tr * w_weight
    bottom = bl * (1 - w_weight) + br * w_weight
    interpolated = top * (1 - h_weight) + bottom * h_weight

    return interpolated

修复点说明 #

1. 用torch.round替代直接int()转换，尽可能保留scale_factor的梯度关联。
2. 用torch.gather替代直接整数索引，确保梯度可传递到输入图像和权重。
3. 所有权重计算使用浮点张量，完整保留梯度流。

梯度传递验证 #

用以下代码验证梯度是否正常传递：

# 测试代码
img = torch.randn(1, 3, 64, 64, requires_grad=True)
scale_factor = torch.tensor(1.5, requires_grad=True)

output = differentiable_interpolate(img, scale_factor)
output.sum().backward()

print("scale_factor的梯度:", scale_factor.grad)  # 输出非零值则说明梯度传递正常

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