应用
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