pytorch 自定义数据集

PyTorch中自定义数据集的步骤如下：

1. 创建一个类，继承自`torch.utils.data.Dataset`。

import torch.utils.data as data

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

    def __getitem__(self, index):
        x = self.data[index]
        y = self.labels[index]
        return x, y

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

2. 实现`__init__`,`__getitem__`和`__len__`方法。其中`__init__`方法初始化数据集，`__getitem__`方法根据索引返回一个样本，`__len__`方法返回数据集的长度。

3. 在训练代码中调用`CustomDataset`类。

from torch.utils.data import DataLoader

train_data = CustomDataset(train_data, train_labels)
train_loader = DataLoader(train_data, batch_size=batch_size, shuffle=True)

4. 使用`DataLoader`类加载数据集。`DataLoader`类可以自动对数据进行批处理、打乱和多进程读取等操作。

以上是自定义数据集的基本步骤，可以根据实际应用场景进一步改进和优化。