torch.Dataset具体使用代码示例

Pytorch框架中，如果要在训练中使用自己定义的数据集，需要继承抽象类Dataset建立新的子集，此时需要重写__len__() 和__getitem__()方法，分别提供数据集的大小，并通过索引号index寻找数据。下面代码提供了如何定义虚拟类以及print出结果。

import os
import torch
import PIL.Image as Image
from torch.utils.data import Dataset

class TensorDataset(Dataset):
    # TensorData继承Dataset，重载了__init__, __getitem__,__lem__
    def __init__(self, data_tensor, target_tensor):
        # 实现将一组Tensor数据对封装成Tensor数据集
        self.data_tensor = data_tensor
        self.target_tensor = target_tensor

    def __getitem__(self, index):
        # 能够通过index得到数据集的数据
        return self.data_tensor[index], self.target_tensor[index]

    def __len__(self):
        # 能够通过len，得到数据集大小
        return self.data_tensor.size(0)  # size(0)返回当前张量维数的第一维

#数据
data_tensor = torch.randn(4, 3)  # 生成四行三列服从正太分布的张量
print(data_tensor)
target_tensor = torch.rand(4)  # 四个元素，服从均匀分布的张量
print(target_tensor)

# 将数据封装成Dataset(用定义的子类TensorDataset类)
tensor_dataset = TensorDataset(data_tensor, target_tensor)
#使用索引调用数据并返回数据大小
print('tensor_data[0]: ', tensor_dataset[0])
print('tensor_data[1]: ', tensor_dataset[1])
print('len os tensor_dataset: ', len(tensor_dataset)) # len返回数据的行数

运行结果：