PyTorch学习笔记(一)自定义数据集

发表于 更新于 分类于 ， Waline： 阅读次数： 本文字数： 1.2k 阅读时长 ≈ 3 分钟

　　利用PyTorch进行深度学习时，首先要做的就是自定义一个数据集类，用于告诉PyTorch如何加载数据集。继承torch.utils.data.Dataset抽象类后，重写两个方法，即可完成自定义数据集的创建：

  • __len__(self)：用于返回数据集的长度
  • __getitem__(self, index)：用于返回第i个数据，传入一个index参数，表示索引

代码示例

　　在这里，数据集大小并不大，因此在构造函数中，直接将所有数据缓存至内存中，内部使用一个列表存储数据。__getitem__(self, index)方法返回2个变量，用于表示第i个label及data，同时一并将data放入Tensor。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
 
import torch
from torch.utils.data import Dataset
from utils import converter


class BallDataset(Dataset):
    """自定义数据集"""

    def __init__(self, data_file):
        """
        :param data_file: 数据集文件路径
        """
        if not os.path.exists(data_file) or not os.path.isfile(data_file):
            raise FileNotFoundError("数据集文件不存在：{}".format(data_file))
        data = []
        file = open(data_file, "r")
        for line in file:
            row = line.split("\t")
            data.append({"label": int(row[0]), "data": row[1:8]})
        file.close()
        self.data = data

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

    def __getitem__(self, index):
        item = self.data[index]
        return item["label"], torch.IntTensor(item["data"])

参考文献

扩展阅读