在 PyTorch 中创建自定义数据集

Question

问题

在 PyTorch 中，我正在尝试编写一个类，它可以使用 dataset.data 和 dataset.label 等语法分别返回整个 data 和 label。代码骨架如下所示：

class MyDataset(object):
  data = _get_data()
  label = _get_label()
  def __init__(self, dir, transforms):
    self.img_list = ... # all image paths loaded from dir
    # do something 

  def __getitem__(self):
    # do something
    return data, label

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

  def _get_data():
    # do something

  def _get_label():
    # do something

但是，当我使用dataset.data 和dataset.label 访问相应的变量时，没有返回任何内容。

我想知道为什么会出现这种情况以及如何解决这个问题。

编辑

感谢大家的关注。

我自己解决了这个问题。解决方案非常简单，它只是利用了类变量的属性。

class FaceDataset(object):
    # class variable
    data = None
    label = None

    def __init__(self, root, transforms=None):
        # read img_list from root
        img_list = ...
        self.transforms = ...
        FaceDataset.data = FaceDataset._get_data(self.img_list, self.transforms)
        FaceDataset.label = FaceDataset._get_label(self.img_list)

    @classmethod
    def _get_data(cls, img_list, transforms):
        data_list = []
        for img_path in img_list:
            data_list.append(transforms(Image.open(img_path)).unsqueeze(0))
        return torch.stack(data_list, dim=0)

    @classmethod
    def _get_label(cls, img_list):
        label = torch.zeros(len(img_list))
        for i, img_path in enumerate(img_list):
            label[i] = ...
        return label

    def __getitem__(self, index):
        img_path = self.img_list[index]
        label = ...

        # read image from file
        data = Image.open(img_path)
        # apply transform defined in __init__
        data = self.transforms(data)

        return data, label

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

Answer 1

在 Python 中创建自定义数据集的“正常”方法已经在 SO 上得到了回答 here。恰好有一个官方 PyTorch tutorial 用于此。

举个简单的例子，你可以阅读 PyTorch MNIST 数据集代码here（此数据集在此 PyTorch example code 中用于进一步说明）。最后，您可以在此 torchvision datasets list 中找到其他数据集实现（单击数据集名称，然后单击数据集文档中的“源”按钮，以访问数据集的 PyTorch 实现）。