如何在 Pytorch 中创建自定义数据加载器？

Question

我有一个包含图像路径的文件，我想加载到 Pytorch 中，同时利用内置的数据加载器功能（多进程加载管道、数据扩充等）。def create_links():    data_dir = "/myfolder"    full_path_list = []    assert os.path.isdir(data_dir)    for _, _, filenames in os.walk(data_dir):        for filename in filenames:            full_path_list.append(os.path.join(data_dir, filename))    with open(config.data.links_file, 'w+') as links_file:        for full_path in full_path_list:            links_file.write(f"{full_path}\n")def read_links_file_to_list():    config = ConfigProvider.config()    links_file_path = config.data.links_file    if not os.path.isfile(links_file_path):        raise RuntimeError("did you forget to create a file with links to images? Try using 'create_links()'")    with open(links_file_path, 'r') as links_file:        return links_file.readlines()所以我有一个文件列表（或一个生成器，或任何有效的东西）file_list = read_links_file_to_list()，.我如何围绕它构建一个 Pytorch 数据加载器，我将如何使用它？

Answer 1

你想要的是一个Custom Dataset。该__getitem__方法是您应用数据增强等转换的地方。为了让您了解它在实践中的样子，您可以看看我前几天写的这个自定义数据集：

class GTSR43Dataset(Dataset):

    """German Traffic Sign Recognition dataset."""

    def __init__(self, root_dir, train_file, transform=None):

        self.root_dir = root_dir

        self.train_file_path = train_file

        self.label_df = pd.read_csv(os.path.join(self.root_dir, self.train_file_path))

        self.transform = transform

        self.classes = list(self.label_df['ClassId'].unique())

    def __getitem__(self, idx):

        """Return (image, target) after resize and preprocessing."""

        img = os.path.join(self.root_dir, self.label_df.iloc[idx, 7])

        

        X = Image.open(img)

        y = self.class_to_index(self.label_df.iloc[idx, 6])

        if self.transform:

            X = self.transform(X)

        return X, y

    

    def class_to_index(self, class_name):

        """Returns the index of a given class."""

        return self.classes.index(class_name)

    

    def index_to_class(self, class_index):

        """Returns the class of a given index."""

        return self.classes[class_index] 

    

    def get_class_count(self):

        """Return a list of label occurences"""

        cls_count = dict(self.label_df.ClassId.value_counts())

#         cls_percent = list(map(lambda x: (1 - x / sum(cls_count)), cls_count))

        return cls_count

    

    def __len__(self):

        """Returns the length of the dataset."""

        return len(self.label_df)