PyTorch DataLoader为何改变输入维度？如何修复？

问题原因分析 #

你的输出维度异常主要来自两个核心问题：

1. 对DataLoader批次结构的理解偏差
   当数据集每个样本是(data1, data2)（每个元素维度为[3,128,128]）时，DataLoader默认的collate_fn会将所有样本的data1堆叠成维度为[batch_size, 3, 128, 128]的张量，data2同理。因此迭代DataLoader得到的结果是两个张量组成的元组，而非你代码中假设的(batch, label)结构。

2. 代码中的解构与索引错误

   • 你使用batch, _ = next(data_iter)，但你的数据集样本并没有标签，这里的_是多余的，导致batch实际对应堆叠后的data1张量。
   • 若你的batch_size=1，data1张量的维度是[1,3,128,128]，此时batch[0]会取出第一个样本的data1（维度[3,128,128]），再取batch[0][0]就会得到通道维度的第一个切片（维度[128,128]），这就是你看到的维度丢失现象。

修复方案 #

1. 修正批次解构方式 #

直接将迭代结果解构为对应的数据批次，替代错误的batch, _写法：

data_iter = iter(train_data_loader)
data1_batch, data2_batch = next(data_iter)
# 查看正确的批次维度
print(f"Data1 batch shape: {data1_batch.shape}")
print(f"Data2 batch shape: {data2_batch.shape}")

此时输出应为torch.Size([batch_size, 3, 128, 128])，符合预期。

2. 确认batch_size设置 #

检查你设置的batch_size值，若为1则属于特殊场景，需明确批次张量的维度结构；若为大于1的数值，上述解构后的张量会包含正确的批量维度。

3. 验证数据集返回结构 #

确保你的数据集类__getitem__方法返回的是(data1, data2)这样的元组结构，而非嵌套列表或其他异常结构：

from torch.utils.data import Dataset

class CustomDataset(Dataset):
    def __getitem__(self, idx):
        # 数据加载逻辑示例
        data1 = ...  # shape [3,128,128]
        data2 = ...  # shape [3,128,128]
        return data1, data2  # 返回元组结构

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