pytorch自定义数据集DataLoder

最新推荐文章于 2025-10-25 12:34:27 发布

原创最新推荐文章于 2025-10-25 12:34:27 发布 · 924 阅读

4 ·

CC 4.0 BY-SA版权

文章标签：

#pytorch

pytorch 专栏收录该内容

2 篇文章

订阅专栏

介绍了PyTorch官方例程中数据加载与处理相关内容。torch.utils.data.Dataset是dataset抽象类，可通过继承它定义自己的dataset，需重写len方法返回数据集长度、getitem方法根据索引采样，最后实例化数据集对象。

部署运行你感兴趣的模型镜像

pytorch官方例程：
DATA LOADING AND PROCESSING TUTORIAL
torch.utils.data.Dataset 是dataset的抽象类，我们可以同过继承Dataset来定义自己的dataset，然后重写类里的两个方法：

len 返回数据集的长度
getitem 根据索引对数据集采样

class FaceLandmarksDataset(Dataset):
    """Face Landmarks dataset."""

    def __init__(self, csv_file, root_dir, transform=None):
        """
        Args:
            csv_file (string): Path to the csv file with annotations.
            root_dir (string): Directory with all the images.
            transform (callable, optional): Optional transform to be applied
                on a sample.
        """
        self.landmarks_frame = pd.read_csv(csv_file)
        self.root_dir = root_dir
        self.transform = transform

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

    def __getitem__(self, idx):
        img_name = os.path.join(self.root_dir,
                                self.landmarks_frame.iloc[idx, 0])
        image = io.imread(img_name)
        landmarks = self.landmarks_frame.iloc[idx, 1:].as_matrix()
        landmarks = landmarks.astype('float').reshape(-1, 2)
        sample = {'image': image, 'landmarks': landmarks}

        if self.transform:
            sample = self.transform(sample)

        return sample

实列化数据集对象

face_dataset = FaceLandmarksDataset(csv_file='data/faces/face_landmarks.csv',
                                    root_dir='data/faces/')
dataset_loader = torch.utils.data.DataLoader(face_dataset,
                                             batch_size=4, shuffle=True,
                                             num_workers=4)