PyTorch中DataLoader返回长度为16而非预期1的原因咨询

问题重现 #

用户执行代码：

img_path = 'G:/tiff/NC_H08_20220419_0600.tif'
img = io.imread(img_path).astype(np.float32)
print(img.shape)
data_tf = torchvision.transforms.Compose([torchvision.transforms.ToTensor()])
train_data = data_tf(img)
print(train_data.shape)
train_loader = DataLoader(dataset=train_data, batch_size=1)
print(len(train_loader))

得到输出：

(2486, 2755, 16)
torch.Size([16, 2486, 2755])
16

用户原本预期len(train_loader)为1，但实际输出是16，以下是原因和解决办法：

原因分析 #

• ToTensor()会将numpy数组的维度从(H, W, C)转换为PyTorch标准的(C, H, W)，原本的16个通道被放到了张量的第一个维度。
• DataLoader默认将输入张量的第一个维度视为样本数量维度，它把这16个通道当成了16个独立样本处理。
• 当batch_size=1时，每个批次仅包含1个样本，总批次数量等于样本数，也就是16，所以len(train_loader)返回16。

解决办法 #

要把整个图像作为1个样本，需要给张量添加一个批次维度，让第一个维度代表样本数（这里是1）：

# 在转换后添加维度
train_data = data_tf(img).unsqueeze(0)
# 此时train_data.shape为torch.Size([1, 16, 2486, 2755])

这样DataLoader会把第一个维度的长度1当作样本总数，len(train_loader)就会返回1，符合预期。

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