pytorch 构建自己的数据集并读取

定义自己的导入数据方法

数据形式可以自定义，例子中的数据形式为[a,b,label]

import torch
import numpy as np
from torch.utils.data import DataLoader

# 定义GetLoader类，继承Dataset方法，并重写__getitem__()和__len__()方法
class GetLoader(torch.utils.data.Dataset):
    # 初始化函数，得到数据
    def __init__(self, data1, data2,data_label):
        self.data = data1
        self.data2 = data2
        self.label = data_label
    # index是根据batchsize划分数据后得到的索引，将batchn内的data及其对应的labels一起返回
    def __getitem__(self, index):
        data = self.data[index]
        data2 = self.data2[index]
        labels = self.label[index]
        return data, data2,labels
    # 该函数返回数据大小长度，目的是DataLoader方便划分，因为data1,data2和label的长度一致，任选一个即可
    def __len__(self):
        return len(self.data)

应用上面的导入数据方法，加载数据

例子中的数据是随机生成的，应用时，按照自己的数据格式进行修改

# 随机生成数据，大小为10 * 2列
source_data = np.random.rand(10, 2)
source_data2 = np.random.rand(10, 2)
# 随机生成标签，大小为10 * 1列
source_label = np.random.randint(0,2,(10, 1))
# 通过GetLoader将数据进行加载，返回Dataset对象，包含data和labels
torch_data = GetLoader(source_data, source_data2 ,source_label)

读取数据（模型训练或测试时使用）

# 读取数据
datas = DataLoader(torch_data, batch_size=5, shuffle=True)
for i, data in enumerate(datas):
	# i表示第几个batch， data表示该batch对应的数据，包含data和对应的labels
    print("第 {} 个Batch \n{}".format(i, data))

上述内容，参考自：Pytorch加载自己的数据集