美文网首页
pytorch的dataset

pytorch的dataset

作者: 山有木紫 | 来源:发表于2018-03-05 18:36 被阅读0次

dataset至少要override2个函数,一个是__len__,另一个是__getitem__
torch.utils.data.Dataset is an abstract class representing a dataset. Your custom dataset should inherit Dataset and override the following methods:
* __len__ so that len(dataset) returns the size of the dataset.
* __getitem__ to support the indexing such that dataset[i] can be used to get iith sample

官方的例子
http://pytorch.org/tutorials/beginner/data_loading_tutorial.html
代码截取

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

相关文章

网友评论

      本文标题:pytorch的dataset

      本文链接:https://www.haomeiwen.com/subject/woohfftx.html