pytorch: 自定义数据集加载

很多网络在数据加载方式

pytorch 的输入流水线的操作顺序是这样的：

    创建一个 Dataset 对象
    创建一个 DataLoader 对象
    不停的 循环 这个 DataLoader 对象

自己感觉：dataset只是负责一张图片以及它的label标签，dataloader负责调用dataset,使得输出batch_size的图片和他们的标签。

dataset = MyDataset()
dataloader = DataLoader(dataset)
num_epoches = 100
for epoch in range(num_epoches):
    for data in dataloader:

其中里面的MyDataset()是继承torch.utils.data.Dataset类，这是一个抽象类，其他所有不管是公开的官方数据集还是自定义数据集都必须继承这个抽象类（比如MNIST数据集），继承这个抽象类的同时必须重写它的两个函数：__len__()  和    __getitem__()。

• __len__()：返回数据集的大小，比如我的数据集有500张图片，那么就返回500
• __getitem__()：返回一张图片

下面使用代码是自定义的数据加载方式。下面的是猫狗数据集里面的

class MyDataset(Dataset):


    def __init__(self, root_dir, transform=None):

        # os.listdir函数读取路径下所以文件的文件名，并组成一个列表并返回

        self.file = os.listdir(root_dir)

        self.root_dir = root_dir

        self.transform = transform

    def __len__(self):
        return len(self.file)  # 返回这给列表的大小

    def __getitem__(self, index):
        # 将传入路径和文件名组成一个新的地址，这个数据就是单个数据的具体地址，方便之后以地址读取该数据
        img_name = os.path.join(self.root_dir, self.file[index])
        # 从文件名中获取标签（我的数据标签在文件名中，比如一张狗的图片名：cat.4.jpg,4是序号，cat是
        # 标签）
        if img_name[13:16] == 'dog':
            label = 0
        else:
            label = 1
        # 根据上面获得的具体地址，读取这张图片
        image = Image.open(img_name)

        # 对图片进行处理
        image = self.transform(image)
        # print(image.shape)

        # numpy的三个维度顺序为：H * W * C
        # 而torch的张量维度顺序：C * H * W ,所以模型要处理它必须转换成torch的形式

        # 返回数据和标签
        return image, label

里面还使用了transform，如下面的代码可以看出来将上面的单个图片进行了预处理。Resize将读取图片的缩放到150。ToTenser()就比较重要了，这个函数的目的就是读取图片像素并且转化为0-1的数字。Normalize作为垫底的一步也很关键，主要用于把图片数据集的数值转化为标准差和均值都为0.5的数据集，这样数据值就从原来的0到1转变为-1到1。

# 数据预处理
transform = transforms.Compose([
    transforms.Resize(size=(150, 150)),
    transforms.ToTensor(),
    transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])
])

猫狗大战中使用自定义数据读取方式代码

#coding=utf-8
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torchvision import datasets, transforms
from torch.utils.data import Dataset  # 首先导入这个抽象类
from PIL import Image
import os


# 设置超参数
NUM_CLASSES = 2
BATCH_SIZE = 50
EPOCHS = 1
DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# 数据预处理
transform = transforms.Compose([
    transforms.Resize(size=(150, 150)),
    transforms.ToTensor(),
    transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])
])



class MyDataset(Dataset):


    def __init__(self, root_dir, transform=None):

        # os.listdir函数读取路径下所以文件的文件名，并组成一个列表并返回

        self.file = os.listdir(root_dir)

        self.root_dir = root_dir

        self.transform = transform

    def __len__(self):
        return len(self.file)  # 返回这给列表的大小

    def __getitem__(self, index):
        # 将传入路径和文件名组成一个新的地址，这个数据就是单个数据的具体地址，方便之后以地址读取该数据
        img_name = os.path.join(self.root_dir, self.file[index])
        # 从文件名中获取标签（我的数据标签在文件名中，比如一张狗的图片名：cat.4.jpg,4是序号，cat是
        # 标签）
        if img_name[13:16] == 'dog':
            label = 0
        else:
            label = 1
        # 根据上面获得的具体地址，读取这张图片
        image = Image.open(img_name)

        # 对图片进行处理
        image = self.transform(image)
        # print(image.shape)

        # numpy的三个维度顺序为：H * W * C
        # 而torch的张量维度顺序：C * H * W ,所以模型要处理它必须转换成torch的形式

        # 返回数据和标签
        return image, label

dataset_train = MyDataset(root_dir='kaggle/train/', transform=transform)
dataset_test = MyDataset(root_dir='kaggle/test1/', transform=transform)

# 导入数据
train_loader = torch.utils.data.DataLoader(dataset_train, batch_size=BATCH_SIZE, shuffle=True)
# test_loader = torch.utils.data.DataLoader(dataset_test, batch_size=BATCH_SIZE, shuffle=True)


# 定义网络
class ConvNet(nn.Module):
    def __init__(self):
        super(ConvNet, self).__init__()
        self.conv1 = nn.Conv2d(3, 32, 3)
        self.max_pool1 = nn.MaxPool2d(2)
        self.conv2 = nn.Conv2d(32, 64, 3)
        self.max_pool2 = nn.MaxPool2d(2)
        self.conv3 = nn.Conv2d(64, 128, 3)
        self.max_pool3 = nn.MaxPool2d(2)
        self.conv4 = nn.Conv2d(128, 128, 3)
        self.max_pool4 = nn.MaxPool2d(2)
        self.fc1 = nn.Linear(6272, 512)
        self.fc2 = nn.Linear(512, 1)

    def forward(self, x):
        in_size = x.size(0)
        x = self.conv1(x)
        x = F.relu(x)
        x = self.max_pool1(x)
        x = self.conv2(x)
        x = F.relu(x)
        x = self.max_pool2(x)
        x = self.conv3(x)
        x = F.relu(x)
        x = self.max_pool3(x)
        x = self.conv4(x)
        x = F.relu(x)
        x = self.max_pool4(x)
        # 展开
        x = x.view(in_size, -1)
        x = self.fc1(x)
        x = F.relu(x)
        x = self.fc2(x)
        x = torch.sigmoid(x)
        return x


# 实例化模型并且移动到GPU
model = ConvNet().to(DEVICE)
# 选择简单暴力的Adam优化器，学习率调低
optimizer = optim.Adam(model.parameters(), lr=1e-4)


# 定义训练过程
def train(model, device, train_loader, optimizer, epoch):
    model.train()
    for batch_idx, (data, target) in enumerate(train_loader):
        # print(data.shape)
        # print(target.shape)
        # print(target)

        data, target = data.to(device), target.to(device).float().reshape(50, 1)
        optimizer.zero_grad()
        output = model(data)
        loss = F.binary_cross_entropy(output, target)
        loss.backward()
        optimizer.step()
        if (batch_idx + 1) % 10 == 0:
            print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
                epoch, (batch_idx + 1) * len(data), len(train_loader.dataset),
                       100. * (batch_idx + 1) / len(train_loader), loss.item()))


# 定义测试过程
def test(model, device, test_loader):
    model.eval()
    test_loss = 0
    correct = 0
    with torch.no_grad():
        for data, target in test_loader:
            # data, target = data.to(device), target.to(device).float().reshape(50, 1)
            data, target = data.to(device), target.to(device)

            output = model(data)
            # print(output)
            test_loss += F.binary_cross_entropy(output, target).item()  # 将一批的损失相加
            pred = torch.tensor([[1] if num[0] >= 0.5 else [0] for num in output]).to(device)
            correct += pred.eq(target.long()).sum().item()
        print('\nTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n'.format(
            test_loss, correct, len(test_loader.dataset),
            100. * correct / len(test_loader.dataset)))

# 训练
for epoch in range(1, EPOCHS + 1):
    train(model, DEVICE, train_loader, optimizer, epoch)
    # test(model, DEVICE, test_loader)

如果不使用自定义数据加载方式

dataset_train = datasets.ImageFolder(root + 'train', transform)

如果使用这个的话，需要将数据文件夹train分成两个子文件夹dog,cat。其实这个就是标签，需要下面的函数

import os
import shutil


def move(root, file_list, out_dog_path, out_cat_path):
    if not os.path.exists(out_dog_path):
        os.mkdir(out_dog_path)
    if not os.path.exists(out_cat_path):
        os.mkdir(out_cat_path)
    for file_name in file_list:
        print(file_name)
        file_path = os.path.join(root, file_name)
        index = file_name[:3]
        # print(index)
        if index == 'dog':
            shutil.move(file_path, out_dog_path)
        else:
            shutil.move(file_path, out_cat_path)

if __name__ == '__main__':
    #train
    root = 'kaggle/train/'
    file_list = os.listdir(root)
    out_dog_path = 'kaggle/train/dog/'
    out_cat_path = 'kaggle/train/cat/'
    move(root, file_list, out_dog_path, out_cat_path)

下面是训练的代码：

#coding=utf-8
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.autograd import Variable
from torchvision import datasets, transforms

# 设置超参数
BATCH_SIZE = 50
EPOCHS = 1
DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# 数据预处理
transform = transforms.Compose([
    transforms.RandomResizedCrop(150),
    transforms.ToTensor(),
    transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])
])

# 读取数据
root = 'kaggle/'
dataset_train = datasets.ImageFolder(root + 'train', transform)
# dataset_test = datasets.ImageFolder(root + 'test1', transform)

# 导入数据
train_loader = torch.utils.data.DataLoader(dataset_train, batch_size=BATCH_SIZE, shuffle=True)
# test_loader = torch.utils.data.DataLoader(dataset_test, batch_size=BATCH_SIZE, shuffle=True)


# 定义网络
class ConvNet(nn.Module):
    def __init__(self):
        super(ConvNet, self).__init__()
        self.conv1 = nn.Conv2d(3, 32, 3)
        self.max_pool1 = nn.MaxPool2d(2)
        self.conv2 = nn.Conv2d(32, 64, 3)
        self.max_pool2 = nn.MaxPool2d(2)
        self.conv3 = nn.Conv2d(64, 128, 3)
        self.max_pool3 = nn.MaxPool2d(2)
        self.conv4 = nn.Conv2d(128, 128, 3)
        self.max_pool4 = nn.MaxPool2d(2)
        self.fc1 = nn.Linear(6272, 512)
        self.fc2 = nn.Linear(512, 1)

    def forward(self, x):
        in_size = x.size(0)
        x = self.conv1(x)
        x = F.relu(x)
        x = self.max_pool1(x)
        x = self.conv2(x)
        x = F.relu(x)
        x = self.max_pool2(x)
        x = self.conv3(x)
        x = F.relu(x)
        x = self.max_pool3(x)
        x = self.conv4(x)
        x = F.relu(x)
        x = self.max_pool4(x)
        # 展开
        x = x.view(in_size, -1)
        x = self.fc1(x)
        x = F.relu(x)
        x = self.fc2(x)
        x = torch.sigmoid(x)
        return x


# 实例化模型并且移动到GPU
model = ConvNet().to(DEVICE)
# 选择简单暴力的Adam优化器，学习率调低
optimizer = optim.Adam(model.parameters(), lr=1e-4)


# 定义训练过程
def train(model, device, train_loader, optimizer, epoch):
    model.train()
    for batch_idx, (data, target) in enumerate(train_loader):
        # data, target = data.to(device), target.to(device).float().reshape(50, 1)
        data, target = data.to(device), target.to(device).float()
        # print(target.size())
        # print(data.size())
        # print(target)
        optimizer.zero_grad()
        output = model(data)
        # print(output)
        loss = F.binary_cross_entropy(output, target)
        loss.backward()
        optimizer.step()
        if (batch_idx + 1) % 10 == 0:
            print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
                epoch, (batch_idx + 1) * len(data), len(train_loader.dataset),
                       100. * (batch_idx + 1) / len(train_loader), loss.item()))


# 定义测试过程
def test(model, device, test_loader):
    model.eval()
    test_loss = 0
    correct = 0
    with torch.no_grad():
        for data, target in test_loader:
            data, target = data.to(device), target.to(device).float().reshape(50, 1)
            output = model(data)
            # print(output)
            test_loss += F.binary_cross_entropy(output, target, reduction='sum').item()  # 将一批的损失相加
            pred = torch.tensor([[1] if num[0] >= 0.5 else [0] for num in output]).to(device)
            correct += pred.eq(target.long()).sum().item()
        print('\nTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n'.format(
            test_loss, correct, len(test_loader.dataset),
            100. * correct / len(test_loader.dataset)))

# 训练
for epoch in range(1, EPOCHS + 1):
    train(model, DEVICE, train_loader, optimizer, epoch)
    # test(model, DEVICE, test_loader)