PyTorch中如何将通道维度重塑为方形邻域像素的技术咨询

你需要将[b, 4*c, h, w]张量转换为[b, c, 2*h, 2*w]，且让4个通道对应2x2的方形邻域（而非reshape的直线排列），有两种可靠实现方式：

手动维度变换 #

核心逻辑是调整维度顺序，让通道拆分出的2x2维度与空间维度对齐后再合并：

import torch

# 模拟输入：b=2, c=3, h=4, w=4，输入形状为[b,4c,h,w]
x = torch.randn(2, 12, 4, 4)

# 1. 将4c通道拆分为c个2x2的通道组
x_reshaped = x.reshape(2, 3, 2, 2, 4, 4)
# 2. 调整维度顺序，把空间维度h/w与2x2的维度对应
x_permuted = x_reshaped.permute(0, 1, 4, 2, 5, 3)
# 3. 合并维度得到目标形状[b,c,2h,2w]
output = x_permuted.reshape(2, 3, 8, 8)

使用官方PixelShuffle层 #

PyTorch内置的PixelShuffle层专门处理这种通道转空间的上采样，代码更简洁且经过优化：

import torch
import torch.nn as nn

# 初始化上采样层，放大倍数设为2
upsampler = nn.PixelShuffle(upscale_factor=2)
# 输入[b,4c,h,w]，直接输出[b,c,2h,2w]
output = upsampler(x)

两种方式效果完全一致，PixelShuffle底层就是手动维度变换的优化实现，推荐优先使用官方层。

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