pytorch torch.utils.data.Dataset

应用

from torch.utils.data import DataLoader, Dataset
import torch

class TensorDataset(Dataset):
    # TensorDataset继承Dataset, 重载了__init__, __getitem__, __len__
    def __init__(self, data_tensor, target_tensor):
        self.data_tensor = data_tensor
        self.target_tensor = target_tensor
    def __getitem__(self, index):
        return self.data_tensor[index], self.target_tensor[index]
    def __len__(self):
        return self.data_tensor.size(0)

# 生成数据
data_tensor = torch.randn(3, 2)
target_tensor = torch.rand(3)
# 将数据封装成Dataset
tensor_dataset = TensorDataset(data_tensor, target_tensor)
# 可使用索引调用数据
tensor_dataset[0] # 等同于(data_tensor[0],target_tensor[0])
# 可返回数据len
len(tensor_dataset)

API

An abstract class representing a Dataset.

All datasets that represent a map from keys to data samples should subclass it. All subclasses should overwrite getitem(), supporting fetching a data sample for a given key. Subclasses could also optionally overwrite len(), which is expected to return the size of the dataset by many Sampler implementations and the default options of DataLoader.

参考：
https://zhuanlan.zhihu.com/p/28200166
https://pytorch.org/docs/master/data.html