如何逆torch.chunk操作按正确顺序重建3D图像原始张量

问题根源 #

你之前的还原操作仅使用view调整维度尺寸，但分块流程中多次通过concat将不同轴拆分得到的块合并到了第0维的块索引维度，view仅会按内存存储顺序重解释张量形状，不会移动元素位置，因此最终得到的张量元素顺序必然和原始张量不符。

可直接运行的还原代码 #

你的分块逻辑是将形状为[4,4,4]的3D张量在深度、高度、宽度三个轴上各拆分为2个无重叠块，最终得到形状为[8,2,2,2]的分块张量，其中第0维的8个块按「深度分块→高度分块→宽度分块」的顺序排列，每个块尺寸为[2,2,2]。还原时需要先拆分块索引维度，再通过维度置换对齐分块索引和块内维度，最后合并维度得到原始张量：

import torch

# 复现原始分块流程
tens = torch.tensor(range(64)).view(4,4,4)
chunked = torch.chunk(tens, 2, -1)
chunked = torch.stack(chunked)
chunked = torch.chunk(chunked, 2, -2)
chunked = torch.concat(chunked)
chunked = torch.chunk(chunked, 2, -3)
chunked = torch.concat(chunked)

# 还原流程
# 1. 将第0维的块索引拆分为三个轴对应的分块索引，维度形状变为[深度分块数, 高度分块数, 宽度分块数, 块内深度, 块内高度, 块内宽度]
restored = chunked.view(2, 2, 2, 2, 2, 2)
# 2. 置换维度顺序，让每个轴的分块索引和对应块内维度相邻
restored = restored.permute(0, 3, 1, 4, 2, 5)
# 3. 合并相邻维度，回到原始[4,4,4]形状
restored = restored.reshape(4, 4, 4)

# 验证还原正确性
print(torch.allclose(restored, tens))  # 输出True

适配自定义分块大小的方法 #

如果你后续调整分块参数，比如将每个维度拆分为n块（要求原张量各维度尺寸能被n整除），只需要对应修改代码里的分块数和块内尺寸即可：

• 拆分块索引维度时，view的前三个参数填每个轴的分块数，后三个参数填每个块在对应轴上的尺寸
• 维度置换的顺序固定为(0,3,1,4,2,5)，作用是把分块索引和对应块内维度配对
• 最后reshape的参数填原始3D张量的尺寸即可

注意：不要全程仅使用view做还原操作，view不会修改元素在内存中的排列顺序，必须通过permute调整维度位置才能对齐原始张量的元素排列逻辑。

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