pytorch 为什么要按维度加载Dataloader？

Question

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))

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

I think len(train_loader) is 1，but now it is 16, I wonder why.

Answer 1

The DataLoader assumes you pass in a dataset, which is usually not a single piece of data. Therefore, it will interpret the first dimension usually as the batch dimension. So, in your case, it assumes you have 16 pieces of 2D data.

To solve it, add a batch dimension to your train_data . (Or make a Dataset , but that seems like a hassle for your simple use case)