Pytorch - 一文搞懂如何使用Pytorch构建与训练自定义深度学习网络（数据集自定义与加载，模型训练，模型测试，模型保存与加载）

最新推荐文章于 2024-02-26 19:22:31 发布

HW140701

最新推荐文章于 2024-02-26 19:22:31 发布

阅读量3.3k

点赞数 7

分类专栏： Pytorch 文章标签： pytorch 深度学习网络自定义

本文为博主原创文章，未经博主允许不得转载，如需转载请先得到博主的同意，如需疑问，请联系stubbornhuang@qq.com，也可以加入计算机图形图像群526867211，以及访问我的个人站点：www.stubbornhuang.com，谢谢。

本文链接：https://blog.csdn.net/HW140701/article/details/122090966

版权

Pytorch 专栏收录该内容

6 篇文章 4 订阅

订阅专栏

使用Pytorch从零训练一个深度学习网络模型，常需要经过以下步骤：自定义数据集、加载自定义数据集、网络模型结构定义、定义损失函数、定义优化器、训练模型、测试模型、保存与加载模型等步骤。下文将详细阐述如何从零构建并训练一个深度学习网络模型的必要步骤。

1 自定义数据集

from torch.utils.data import Dataset

class torch.utils.data.Dataset

表示数据集Dataset的抽象类，所有其他自定义的数据集都应该继承Dataset类，并且强制重写_len_和_getitem_方法，

from torch.utils.data import Dataset

class MyDataset(Dataset):
    def __init__(self):
        # TODO
        # 1. Initialize file path or list of file names.
        pass
    def __len__(self):
        # TODO
        # 1. Read one data from file (e.g. using numpy.fromfile, PIL.Image.open).
        # 2. Preprocess the data (e.g. torchvision.Transform).
        # 3. Return a data pair (e.g. image and label).
        pass
   	 def __getitem__(self, index):
       # TODO
       # You should change 0 to the total size of your dataset.
       pass

其中：

init() : 将一些初始化过程写在此函数中，比如从文件中加载数据。此构造函数可以自定义参数，比如可以传递数据文件路径以及标签文件路径：

def __init__(self,data_path,label_path):
    pass

传递上述路径之后可以从上述路径中读取数据和对应标签，之后用于模型训练。当然也可在构造函数中添加需要在数据集加载中需要使用的参数，比如指定数据是否需要归一化，是否需要混淆，是否需要镜像等等。

len()：返回所有数据的数量
getitem(self, index)：返回指定索引index的数据与对应标签，另外可在此函数中做数据标准化、数据混淆和数据增加的处理。

形如：

import matplotlib.pyplot as plt
import numpy as np
import torch
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms
from PIL import Image
import os
image_transform = transforms.Compose([
    transforms.Resize(256),               # 把图片resize为256*256
    transforms.RandomCrop(224),           # 随机裁剪224*224
    transforms.RandomHorizontalFlip(),    # 水平翻转
    transforms.ToTensor(),                # 将图像转为Tensor
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])   # 标准化
])
 
class DogVsCatDataset(Dataset):   # 创建一个叫做DogVsCatDataset的Dataset，继承自父类torch.utils.data.Dataset
    def __init__(self, root_dir, train=True, transform=None):
        """
        Args:
            root_dir (string): Directory with all the images.
            transform (callable, optional): Optional transform to be applied on a sample.
        """
        self.root_dir = root_dir
        self.img_path = os.listdir(self.root_dir)
        if train:
            self.img_path = list(filter(lambda x: int(x.split('.')[1]) < 10000, self.img_path))    # 划分训练集和验证集
        else:
            self.img_path = list(filter(lambda x: int(x.split('.')[1]) >= 10000, self.img_path))
        self.transform = transform
 
    def __len__(self):
        return len(self.img_path)
 
    def __getitem__(self, idx):
        image = Image.open(os.path.join(self.root_dir, self.img_path[idx]))
        label = 0 if self.img_path[idx].split('.')[0] == 'cat' else 1        # label, 猫为0，狗为1
        if self.transform:
            image = self.transform(image)
        label = torch.from_numpy(np.array([label]))
        return image, label
 
if __name__ == '__main__':
    catanddog_dataset = DogVsCatDataset(root_dir='/Users/wangpeng/Desktop/train',
                                        train=False,
                                        transform=image_transform)
    train_loader = DataLoader(catanddog_dataset, batch_size=8, shuffle=True, num_workers=4)   # num_workers=4表示用4个线程读取数据
    image, label = iter(train_loader).next()   # iter()函数把train_loader变为迭代器，然后调用迭代器的next()方法
    sample = image[0].squeeze()
    sample = sample.permute((1, 2, 0)).numpy()
    sample *= [0.229, 0.224, 0.225]
    sample += [0.485, 0.456, 0.406]
    sample = np.clip(sample, 0, 1)
    plt.imshow(sample)
    plt.show()
    print('Label is: {}'.format(label[0].numpy()))

自定义数据集也可参考：

2 加载自定义数据集

from torch.utils.data import DataLoader

class torch.utils.data.DataLoader(dataset, batch_size=1, shuffle=False, sampler=None, num_workers=0, collate_fn=<function default_collate>, pin_memory=False, drop_last=False)

参数：

dataset (Dataset) – 加载数据的数据集。这个参数可以传递自定义的数据集对象
batch_size (int, optional) – 每个batch加载多少个样本(默认: 1)。
shuffle (bool, optional) – 设置为True时会在每个epoch重新打乱数据(默认: False).
sampler (Sampler, optional) – 定义从数据集中提取样本的策略。如果指定，则忽略shuffle参数。
num_workers (int, optional) – 用多少个子进程加载数据。0表示数据将在主进程中加载(默认: 0)
collate_fn (callable, optional) – 将一个list的sample组成一个mini-batch的函数
pin_memory (bool, optional) – 如果设置为True，那么data loader将会在返回它们之前，将tensors拷贝到CUDA中的固定内存（CUDA pinned memory）中.
drop_last (bool, optional) – 如果数据集大小不能被batch size整除，则设置为True后可删除最后一个不完整的batch。如果设为False并且数据集的大小不能被batch size整除，则最后一个batch将更小。(默认: False)。这个是对最后的未完成的batch来说的，比如你的batch_size设置为64，而一个epoch只有100个样本，那么训练的时候后面的36个就被扔掉了，如果为False（默认），那么会继续正常执行，只是最后的batch_size会小一点。
timeout(numeric, optional): 如果是正数，表明等待从worker进程中收集一个batch等待的时间，若超出设定的时间还没有收集到，那就不收集这个内容了。这个numeric应总是大于等于0。默认为0
worker_init_fn (callable, optional): 每个worker初始化函数 If not None, this will be called on each

一般的，加载自定义数据集的方法为：

mydataset = MyDataset()

dataloader = DataLoader(dataset=mydataset,batch_size=32,shuffle=True,num_workers=8,drop_last=True)

for batch_index ,(data,label) in enumerate(dataloader):
    data = Variable(data.cuda(0),requires_grad=False)
    label = Variable(label.cuda(0),requires_grad=False)

3 定义模型

class torch.nn.Module

所有网络模型的基类，所有自定义的网络模型都需要继承此类，并且重写forward方法，forword定义了每次执行的计算步骤。一般在自定义的网络模型的构造函数_init_中注册网络模型需要使用各种神经网络层，然后在forward函数中使用在构造函数设定的各种神经网络层进行组合，构建自定义的神经网络。

一个简单的网络模型定义如下：

import torch.nn as nn
import torch.nn.functional as F

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(1, 20, 5, 1)
        self.conv2 = nn.Conv2d(20, 50, 5, 1)
        self.fc1 = nn.Linear(4 * 4 * 50, 500)
        self.fc2 = nn.Linear(500, 10)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        x = F.max_pool2d(x, 2, 2)
        x = F.relu(self.conv2(x))
        x = F.max_pool2d(x, 2, 2)
        x = x.view(-1, 4 * 4 * 50)
        x = F.relu(self.fc1(x))
        x = self.fc2(x)
        return F.log_softmax(x, dim=1)

class torch.nn.Module常用的成员函数方法：

cpu(device_id=None) ：将所有的模型参数(parameters)和buffers复制到CPU
cuda(device_id=None)：将所有的模型参数(parameters)和buffers赋值GPU，device_id如果指定的话，所有的模型参数都会复制到指定的设备上
double()：将parameters和buffers的数据类型转换成double。
float()：将parameters和buffers的数据类型转换成float。
half()：将parameters和buffers的数据类型转换成half。
eval()：将模型设置成evaluation模式，仅仅当模型中有Dropout和BatchNorm是才会有影响。
train(mode=True)：将module设置为 training mode。仅仅当模型中有Dropout和BatchNorm是才会有影响。
zero_grad()：将module中的所有模型参数的梯度设置为0。

4 定义损失函数

在模型网络定义好之后，我们需要定义在模型训练过程中需要使用到的损失函数，形如：

loss = nn.CrossEntropyLoss().cuda(output_device)

pytorch中定义了很多损失函数，如：

torch.nn.L1Loss(size_average=True)：创建一个衡量输入x(模型预测输出)和目标y之间差的绝对值的平均值的标准。
torch.nn.MSELoss(size_average=True)：创建一个衡量输入x(模型预测输出)和目标y之间均方误差标准。
torch.nn.CrossEntropyLoss(weight=None, size_average=True)：此标准将LogSoftMax和NLLLoss集成到一个类中，当训练一个多类分类器的时候，这个方法是十分有用的。
torch.nn.NLLLoss(weight=None, size_average=True)：负的log likelihood loss损失。用于训练一个n类分类器。
torch.nn.NLLLoss2d(weight=None, size_average=True)：对于图片的 negative log likehood loss。计算每个像素的 NLL loss。
torch.nn.KLDivLoss(weight=None, size_average=True)：计算 KL 散度损失。KL散度常用来描述两个分布的距离，并在输出分布的空间上执行直接回归是有用的。
torch.nn.BCELoss(weight=None, size_average=True)：计算 target 与 output 之间的二进制交叉熵。
torch.nn.MarginRankingLoss(margin=0, size_average=True)
torch.nn.HingeEmbeddingLoss(size_average=True)：这个loss通常用来测量两个输入是否相似，即：使用L1 成对距离。典型是用在学习非线性 embedding或者半监督学习中。
torch.nn.MultiLabelMarginLoss(size_average=True)：计算多标签分类的 hinge loss(margin-based loss) 。
torch.nn.SmoothL1Loss(size_average=True)：平滑版L1 loss。
torch.nn.SoftMarginLoss(size_average=True)：创建一个标准，用来优化2分类的logistic loss。
torch.nn.MultiLabelSoftMarginLoss(weight=None, size_average=True)：创建一个标准，基于输入x和目标y的 max-entropy，优化多标签 one-versus-all 的损失。
torch.nn.CosineEmbeddingLoss(margin=0, size_average=True)：此标准使用cosine距离测量两个输入是否相似，一般用来用来学习非线性embedding或者半监督学习。
torch.nn.MultiMarginLoss(p=1, margin=1, weight=None, size_average=True)：用来计算multi-class classification的hinge loss（magin-based loss）。

通常需要根据需要具体的任务类型选择合适的损失函数，比如当面临一个多分类任务时，这时候选择torch.nn.CrossEntropyLoss作为损失函数就是最合适的。

5 定义优化器

torch.optim

torch.optim是一个实现了各种优化算法的库，用于保存当前模型参数状态并基于所计算的得到的梯度进行模型参数更新。我们可以设置optimizer的参数，比如学习率，权重衰减等等。

在pytorch中定义优化器可参考：

optimizer = optim.SGD(model.parameters(), lr = 0.01, momentum=0.9)
optimizer = optim.Adam([var1, var2], lr = 0.0001)

所有的optimizer都实现了step()方法，这个方法会更新所有的参数。一旦梯度被如backward()之类的函数计算好后，我们就可以调用这个函数。

单次训练过程中使用优化器进行参数更新可参考如下：

# 定义损失函数
loss_fn = nn.CrossEntropyLoss().cuda(output_device)

# 定义优化器
optimizer = optim.Adam([var1, var2], lr = 0.0001)

for batch_index ,(data,label) in enumerate(dataloader):
    data = Variable(data.cuda(0),requires_grad=False)
    label = Variable(label.cuda(0),requires_grad=False)
    
    # 清空所有被优化过的梯度
    optimizer.zero_grad()
    
    # 执行模型forward
    pred_label = model(data)
    
    # 损失函数比较差异
    loss = loss_fn(pred_label, label)
    
    # 反向传播
    loss.backward()
    
    # 优化模型参数
    optimizer.step()

6 训练模型

6.1 模型训练前的准备步骤

自定义数据集
定义网络模型
定义损失函数
定义优化器

具体步骤可参考上述1 - 5节的内容。

6.2 模型训练的必要步骤

加载数据集
创建模型对象
定义训练次数即epoch
在每一个epoch中对模型进行训练

在单个epoch中，通常需要使用model.train()先将模型切换到train状态，然后使用dataloader加载数据集并遍历数据以及对应标签，然后使用optimizer.zero_grad()清空所有优化过的梯度，然后将使用将得到的当前遍历的数据输入到模型中得到预测的结果，使用定义的损失函数计算预测结果与真实标签的loss，使用该loss进行反向传播，然后使用optimizer.step()基于所计算的得到的梯度进行模型参数更新。

单次的训练过程示例代码形如：

def train(args, model, device, train_loader, optimizer, epoch):
    # 切换到train状态
    model.train()
    for batch_idx, (data, target) in enumerate(train_loader):
        # 从数据集获取数据以及对应标签
        data, target = data.to(device), target.to(device)
        # 清空所有优化过的梯度
        optimizer.zero_grad()
        # 执行模型forward
        output = model(data)
        # 计算损失
        loss = torch.nn.functional.nll_loss(output, target)
        # 反向传播
        loss.backward()
        # 优化模型参数
        optimizer.step()

7 测试模型

在模型测试的过程中，需要先将model.eval将模型设置成evaluation模式，然后使用with torch.no_grad()包裹测试代码，被with torch.no_grad()包裹的代码不需要跟踪反向梯度计算。然后使用dataloader加载数据集并遍历数据以及对应标签，然后将使用将得到的当前遍历的数据输入到模型中得到预测的结果。

单次的测试模型的示例代码如下：

def test(args, model, device, test_loader):
    # 切换到evaluation`模式
    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)
            # 执行模型forward
            output = model(data)
            test_loss += torch.nn.functional.nll_loss(output, target, reduction='sum').item()  # sum up batch loss
            # 获取推理的最大可能标签
            pred = output.argmax(dim=1, keepdim=True)  # get the index of the max log-probability
            correct += pred.eq(target.view_as(pred)).sum().item()

8 保存与加载模型

8.1 保存和加载整个模型

保存整个模型：

torch.save(model_object, 'model.pt')

加载整个模型：

model = torch.load('model.pt')

8.2 仅保存和加载模型参数

保存模型参数：

torch.save(model_object.state_dict(), 'model.pt')

加载模型参数：

model_object.load_state_dict(torch.load('model.pt'))

9 pytorch CNN 代码实战说明

import argparse
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torchvision import datasets, transforms

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(1, 20, 5, 1)
        self.conv2 = nn.Conv2d(20, 50, 5, 1)
        self.fc1 = nn.Linear(4 * 4 * 50, 500)
        self.fc2 = nn.Linear(500, 10)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        x = F.max_pool2d(x, 2, 2)
        x = F.relu(self.conv2(x))
        x = F.max_pool2d(x, 2, 2)
        x = x.view(-1, 4 * 4 * 50)
        x = F.relu(self.fc1(x))
        x = self.fc2(x)
        return F.log_softmax(x, dim=1)


def train(args, 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)
        optimizer.zero_grad()
        output = model(data)
        loss = F.nll_loss(output, target)
        loss.backward()
        optimizer.step()
        if batch_idx % args.log_interval == 0:
            print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
                epoch, batch_idx * len(data), len(train_loader.dataset),
                       100. * batch_idx / len(train_loader), loss.item()))


def test(args, 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)
            output = model(data)
            test_loss += F.nll_loss(output, target, reduction='sum').item()  # sum up batch loss
            pred = output.argmax(dim=1, keepdim=True)  # get the index of the max log-probability
            correct += pred.eq(target.view_as(pred)).sum().item()

    test_loss /= len(test_loader.dataset)

    print('\nTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n'.format(
        test_loss, correct, len(test_loader.dataset),
        100. * correct / len(test_loader.dataset)))


def main():
    parser = argparse.ArgumentParser(description='PyTorch MNIST Example')
    parser.add_argument('--batch-size', type=int, default=64, metavar='N',
                        help='input batch size for training (default: 64)')
    parser.add_argument('--test-batch-size', type=int, default=1000, metavar='N',
                        help='input batch size for testing (default: 1000)')
    parser.add_argument('--epochs', type=int, default=10, metavar='N',
                        help='number of epochs to train (default: 10)')
    parser.add_argument('--lr', type=float, default=0.01, metavar='LR',
                        help='learning rate (default: 0.01)')
    parser.add_argument('--momentum', type=float, default=0.5, metavar='M',
                        help='SGD momentum (default: 0.5)')
    parser.add_argument('--no-cuda', action='store_true', default=False,
                        help='disables CUDA training')
    parser.add_argument('--seed', type=int, default=1, metavar='S',
                        help='random seed (default: 1)')
    parser.add_argument('--log-interval', type=int, default=10, metavar='N',
                        help='how many batches to wait before logging training status')
    parser.add_argument('--save-model', action='store_true', default=False,
                        help='For Saving the current Model')

    args = parser.parse_args()
    use_cuda = not args.no_cuda and torch.cuda.is_available()
    torch.manual_seed(args.seed)
    device = torch.device("cuda:0" if use_cuda and torch.cuda.is_available() else "cpu")
    print(device)

    kwargs = {'num_workers': 4, 'pin_memory': True} if use_cuda else {}
    train_loader = torch.utils.data.DataLoader(
        datasets.MNIST('./mnist', train=True, download=True,
                       transform=transforms.Compose([
                           transforms.ToTensor(),
                           transforms.Normalize((0.1307,), (0.3081,))
                       ])),
        batch_size=args.batch_size, shuffle=True, **kwargs)
    test_loader = torch.utils.data.DataLoader(
        datasets.MNIST('./mnist', train=False, transform=transforms.Compose([
            transforms.ToTensor(),
            transforms.Normalize((0.1307,), (0.3081,))
        ])),
        batch_size=args.test_batch_size, shuffle=True, **kwargs)

    model = Net().to(device)
    optimizer = optim.SGD(model.parameters(), lr=args.lr, momentum=args.momentum)

    for epoch in range(1, args.epochs + 1):
        train(args, model, device, train_loader, optimizer, epoch)

    test(args, model, device, test_loader)

    if (args.save_model):
        torch.save(model.state_dict(), "mnist_cnn.pt")

if __name__ == '__main__':
    main()

9.1 定义模型

使用下列代码定义了简单的CNN模型：

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(1, 20, 5, 1)
        self.conv2 = nn.Conv2d(20, 50, 5, 1)
        self.fc1 = nn.Linear(4 * 4 * 50, 500)
        self.fc2 = nn.Linear(500, 10)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        x = F.max_pool2d(x, 2, 2)
        x = F.relu(self.conv2(x))
        x = F.max_pool2d(x, 2, 2)
        x = x.view(-1, 4 * 4 * 50)
        x = F.relu(self.fc1(x))
        x = self.fc2(x)
        return F.log_softmax(x, dim=1)

9.2 数据集加载

在上述代码中，首先使用了pytorch自带的MNIST数据集作为数据集，并使用了DataLoader进行了数据集加载

    train_loader = torch.utils.data.DataLoader(
        datasets.MNIST('./mnist', train=True, download=True,
                       transform=transforms.Compose([
                           transforms.ToTensor(),
                           transforms.Normalize((0.1307,), (0.3081,))
                       ])),
        batch_size=args.batch_size, shuffle=True, **kwargs)
    test_loader = torch.utils.data.DataLoader(
        datasets.MNIST('./mnist', train=False, transform=transforms.Compose([
            transforms.ToTensor(),
            transforms.Normalize((0.1307,), (0.3081,))
        ])),
        batch_size=args.test_batch_size, shuffle=True, **kwargs)

9.3 定义损失函数

import torch.nn.functional as F

9.4 定义优化器

optimizer = optim.SGD(model.parameters(), lr=args.lr, momentum=args.momentum)

9.5 模型训练

单次epoch模型训练：

def train(args, 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)
        optimizer.zero_grad()
        output = model(data)
        loss = F.nll_loss(output, target)
        loss.backward()
        optimizer.step()
        if batch_idx % args.log_interval == 0:
            print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
                epoch, batch_idx * len(data), len(train_loader.dataset),
                       100. * batch_idx / len(train_loader), loss.item()))

多次epoch模型训练：

    for epoch in range(1, args.epochs + 1):
        train(args, model, device, train_loader, optimizer, epoch)

9.6 模型测试

def test(args, 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)
            output = model(data)
            test_loss += F.nll_loss(output, target, reduction='sum').item()  # sum up batch loss
            pred = output.argmax(dim=1, keepdim=True)  # get the index of the max log-probability
            correct += pred.eq(target.view_as(pred)).sum().item()

    test_loss /= len(test_loader.dataset)

    print('\nTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n'.format(
        test_loss, correct, len(test_loader.dataset),
        100. * correct / len(test_loader.dataset)))

9.7 模型保存

    if (args.save_model):
        torch.save(model.state_dict(), "mnist_cnn.pt")

参考链接

有兴趣可以访问我的个站：https://www.stubbornhuang.com/