DataLoader 中的批量大小答案

【问题标题】：Batchsize in DataLoaderDataLoader 中的批量大小
【发布时间】：2021-10-28 04:30:14
【问题描述】：

我有两个张量：

x[train], y[train]

形状是

(311, 3, 224, 224), (311) # 311 Has No Information

我想用DataLoader来批量加载，我写的代码是：

from torch.utils.data import Dataset

class KD_Train(Dataset):

    def __init__(self,a,b):
        self.imgs = a
        self.index = b

    def __len__(self):
        return len(self.imgs)

    def __getitem__(self,index):
        return self.imgs, self.index

kdt = KD_Train(x[train], y[train])

train_data_loader = Data.DataLoader(
    kdt,
    batch_size = 64,
    shuffle = True,
    num_workers = 0)

for step, (a,b) in enumerate (train_data_loader):
    print(a.shape)
    break

但它显示：

(64, 311, 3, 224, 224)

DataLoader只是直接添加维度，而不是选择一些批次，有人知道我该怎么办吗？

【问题讨论】：

标签： pytorch dataloader

【解决方案1】：

您的数据集的 __getitem__ 方法应返回单个元素：

def __getitem__(self, index):
    return self.imgs[index], self.index[index]

【讨论】：