Torch如何自定义一个Dataset类-开发者知识库平台

Torch如何自定义一个Dataset类

Torch

1227

2024/2/28 11:53:10

栏目: 深度学习

要自定义一个Dataset类，可以继承自torch.utils.data.Dataset，并实现其中的__len__和__getitem__方法来定义数据集的长度和获取数据的方式。

下面是一个简单的例子：

import torch
from torch.utils.data import Dataset

class CustomDataset(Dataset):
    def __init__(self, data):
        self.data = data

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

    def __getitem__(self, idx):
        sample = self.data[idx]
        return sample

# 创建一个数据集实例
data = [1, 2, 3, 4, 5]
custom_dataset = CustomDataset(data)

# 获取数据集的长度
print(len(custom_dataset))

# 获取数据集中第一个样本
print(custom_dataset[0])

在上面的例子中，我们定义了一个CustomDataset类，它接受一个数据列表作为输入，并实现了__len__方法和__getitem__方法。通过实例化CustomDataset类，我们可以获取数据集的长度并获取数据集中的样本。

辰迅云「云服务器」，即开即用、新一代英特尔至强铂金CPU、三副本存储NVMe SSD云盘，价格低至29元/月。点击查看>>

Torch如何自定义一个Dataset类

最新知识库

相关标签