TensorDataset为何拆分数据为小批量？示例疑问与内置方案咨询

为什么遍历出来的batch长度是1？ #

TensorDataset的核心逻辑是按第一维度对齐样本：

• 它要求所有输入张量的第一维度长度一致，这个维度代表「样本数」
• 遍历Dataset时，每个元素是一个元组，元组的长度等于你传入的张量个数，每个元素对应各张量在该样本位置的切片

你传入了一个形状为(20, 3)的张量，TensorDataset会把它理解为「20个样本，每个样本有3个特征」。所以遍历出来的每个元素（单个样本）是单元素元组（因为只传了一个张量），元组里的内容是该样本的3个特征（形状(3,)）——这里的batch其实是单个样本，而非批量数据，这就是len(batch)=1的原因。

如何正确获取批量数据？ #

不需要手动拆包，用DataLoader就能实现批量处理：
DataLoader会自动把多个样本拼接成批量张量，你只需要通过batch_size参数指定每次取多少个样本即可。示例代码：

from torch.utils.data import TensorDataset, DataLoader
import torch

data = torch.randint(0, 100, (20, 3), dtype=torch.int32)
tensor_dataset = TensorDataset(data)
# 指定每次取5个样本组成批量
dataloader = DataLoader(tensor_dataset, batch_size=5)

for batch in dataloader:
    print(len(batch))  # 输出：1（元组长度对应传入的张量个数）
    print(batch[0].shape)  # 输出：torch.Size([5, 3])（5个样本，每个3个特征）

如果想把特征维度当作样本数怎么办？ #

如果你实际想处理的是「3个样本，每个样本有20个特征」，只需要转置原始张量即可：

data = data.T  # 形状从(20,3)变为(3,20)
tensor_dataset = TensorDataset(data)
print(len(tensor_dataset))  # 输出：3
for batch in tensor_dataset:
    print(len(batch))  # 输出：1
    print(batch[0].shape)  # 输出：torch.Size([20])

内容的提问来源于stack exchange，提问作者J. Doe