PyTorch中如何对N>3的(N+2)维张量实现N维线性插值？

PyTorch的F.interpolate仅支持最多3个空间维度的插值操作（对应3D/4D/5D张量），要处理N>3的(N+2)维张量（前两维为批量B和通道C），可以利用线性插值的可分离性——N维线性插值等价于在每个空间维度依次执行1D线性插值。

实现代码 #

import torch
import torch.nn.functional as F

def nd_linear_interpolate(x, output_size):
    # 输入张量形状：(B, C, *spatial_dims)
    # output_size：目标空间维度的尺寸元组，长度需与输入的空间维度数一致
    assert len(output_size) == x.dim() - 2, "目标尺寸的长度必须匹配输入的空间维度数量"
    
    y = x
    spatial_dims_count = len(output_size)
    for dim_idx in range(spatial_dims_count):
        # 构造当前插值的目标尺寸：仅修改当前处理的空间维度，其余保持原尺寸
        current_target_size = list(y.shape[2:])
        current_target_size[dim_idx] = output_size[dim_idx]
        # 执行1D线性插值
        y = F.interpolate(y, size=current_target_size, mode="linear", align_corners=False)
    return y

# 测试示例
if __name__ == "__main__":
    x = torch.randn(1, 1, 2, 3, 4, 5, 6, 7)  # 1个批量、1个通道、6个空间维度
    output_size = (7, 6, 5, 4, 3, 2)
    y = nd_linear_interpolate(x, output_size)
    print(f"输入形状: {x.shape}")
    print(f"输出形状: {y.shape}")  # 输出应为 (1, 1, 7, 6, 5, 4, 3, 2)

关键说明 #

• 逐维度迭代：每次仅对一个空间维度执行1D线性插值，其余维度保持原尺寸，循环完成所有空间维度的缩放。
• 等价性：可分离的线性插值组合结果与直接N维线性插值完全一致，符合数学定义。
• 参数调整：align_corners参数可根据需求设置为True或False，默认False与PyTorch内置插值的默认行为对齐。

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