保留梯度的PyTorch张量无卷积快速转换方法

核心方案 #

方法1：基于Unfold的实现 #

利用PyTorch内置的可微分Unfold操作提取空间块，再调整维度结构，全程保留梯度：

1. 为输入张量添加临时batch维度（适配Unfold的4维输入要求）
2. 用kernel_size=2, stride=2提取所有2x2空间块
3. 将展开的空间位置序列重塑为二维结构，移除临时batch维度得到目标形状

import torch
import torch.nn as nn

def reshape_with_unfold(input_tensor):
    D, M, M = input_tensor.shape
    assert M % 2 == 0, "M必须为偶数"
    
    # 添加临时batch维度，转为(1, D, M, M)
    x = input_tensor.unsqueeze(0)
    # 提取步长为2的2x2块，输出形状(1, D*4, (M//2)^2)
    unfold = nn.Unfold(kernel_size=2, stride=2)
    patches = unfold(x)
    # 重塑空间维度并移除batch维度
    output = patches.view(1, D*4, M//2, M//2).squeeze(0)
    return output

方法2：手动切片拼接（更直观） #

直接提取原张量中四个角落的下采样块，在通道维度拼接，完全避免使用卷积或Unfold：

def reshape_manual(input_tensor):
    D, M, M = input_tensor.shape
    assert M % 2 == 0, "M必须为偶数"
    h, w = M//2, M//2
    
    # 提取四个下采样块
    top_left = input_tensor[:, ::2, ::2]    # (D, h, w)
    top_right = input_tensor[:, ::2, 1::2]  # (D, h, w)
    bottom_left = input_tensor[:, 1::2, ::2]# (D, h, w)
    bottom_right = input_tensor[:, 1::2, 1::2]# (D, h, w)
    
    # 通道维度拼接得到(D*4, h, w)
    output = torch.cat([top_left, top_right, bottom_left, bottom_right], dim=0)
    return output

测试验证 #

# 测试示例：输入(3,4,4) → 输出(12,2,2)
input_tensor = torch.randn(3, 4, 4, requires_grad=True)
print(f"输入形状: {input_tensor.shape}")

# 测试Unfold方法
output_unfold = reshape_with_unfold(input_tensor)
print(f"Unfold方法输出形状: {output_unfold.shape}")
# 验证梯度保留
output_unfold.sum().backward()
print(f"输入梯度存在: {input_tensor.grad is not None}")

# 重置梯度后测试手动方法
input_tensor.grad.zero_()
output_manual = reshape_manual(input_tensor)
print(f"手动方法输出形状: {output_manual.shape}")
output_manual.sum().backward()
print(f"输入梯度存在: {input_tensor.grad is not None}")

维度变换说明 #

两种方法均满足：

• 输入(D, M, M) → 输出(D*4, M//2, M//2)
• 所有操作均为PyTorch可微分操作，梯度完整保留
• 无卷积操作，严格遵循类池化的平铺拼接逻辑

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