pytorch读取数据集基本操作(数值型、图片)

一、数值型数据

import torch
from torch.utils.data import TensorDataset, DataLoader, random_split, Dataset
from sklearn import datasets

# 根据Tensor创建数据集
iris = datasets.load_iris()
Tensor_iris = TensorDataset(torch.tensor(iris.data), torch.tensor(iris.target))

# 分割成训练集、验证集、测试集
num_train = int(len(Tensor_iris) * 0.6)
num_valid = int(len(Tensor_iris) * 0.2)
num_test = len(Tensor_iris) - num_train - num_valid
train, valid, test = random_split(Tensor_iris, [num_train, num_valid, num_test])
print(len(train), len(valid), len(test))  # 90 30 30(个数)

# 使用DataLoader加载数据集
train = DataLoader(train, batch_size=16, shuffle=True)
valid = DataLoader(valid, batch_size=16, shuffle=False)
test = DataLoader(test, batch_size=16, shuffle=False)
print(len(train), len(valid), len(test))  # 6 2 2(批数)

数据集直接用的sklearn的鸢尾花数据集。

多特征数值型数据(如上iris.data)

标签(如上iris.target)

二、图片

(1)分类任务图片：

from torch.utils.data import DataLoader, random_split
from torchvision import transforms, datasets

data_transform = transforms.Compose([transforms.Resize((224, 224)), transforms.ToTensor()])

data = datasets.ImageFolder(root="dir", transform=data_transform)

num_train = int(len(data) * 0.6)
num_valid = int(len(data) * 0.2)
num_test = len(data) - num_train - num_valid
train, valid, test = random_split(data, [num_train, num_valid, num_test])

train = DataLoader(train, batch_size=16, shuffle=True)
valid = DataLoader(valid, batch_size=16, shuffle=False)
test = DataLoader(test, batch_size=16, shuffle=False)

data_transform:根据自己的需求进行数据增强。

图片数据集目录例子：dir-----dog

                                        |----cat

                                    dir的下级目录的名要是要分类的标签名，比如我要分类dog、cat，我就把把文件名改成dog、cat,这两个文件夹下是他俩的所有图片数据。

数据集链接：Cats-vs-Dogs | Kaggle或csdn搜cats vs dogs,有人传了百度网盘。

(2)分割任务图片：

        分割任务图片用torchvision的transform来数据增强比较麻烦，因为涉及到mask也需要进行转变，因此推荐使用albumentations库来进行数据增强。我一般是把数据增强后的图片数据先存在硬盘里（确保之后的数据不变，而不是每次训练都要重新数据增强），再读入内存进行训练，则需要构建自己的数据集。

        数据增强后的图片数据是用numpy存储的，格式为(batch_size, channels, height, width)。

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

class datasets(Dataset):
    def __init__(self, data, label):
        self.length = data.shape[0]
        self.data = data
        self.label = label

    def __getitem__(self, index):
        return self.data[index], self.label[index]

    def __len__(self):
        return self.length

train_data = np.load('XXX')
train_label = np.load('XXX')
train_data = torch.from_numpy(train_data).float()
train_label = torch.from_numpy(train_label).float() #因为这次用的均方误差损失，所以转成浮点数，若为交叉熵损失，则为整型才可以。

train_dataset = datasets(train_data, train_label)
train_loader = DataLoader(train_dataset, batch_size=2, shuffle=True)