Pytorch训练代码框架

前言

网上关于Pytorch的基础训练代码框架有很多，但一直想创建一个自己的。于是从头写了一遍，包含了seed指定，网络搭建，DataSet和DataLoader，loss和优化器的定义和使用，模型的保存和载入。

代码

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

parser = argparse.ArgumentParser(description="输入参数")
parser.add_argument('-batch_size', type=int, default=16, help="BatchSize")
parser.add_argument('-epochs', type=int, default=20, help="Epoch Num")
parser.add_argument('-train', type=str, default='True', help="Train or Eval")
parser.add_argument('-cuda_id', type=int, default='0', help="CUDA序号")

parser.add_argument('-model_save_path', type=str,
                    default='./model', help="model保存路径")
parser.add_argument('-pretrain_model_path', type=str,
                    default='./model/c19.pth', help="预训练模型路径")

args = parser.parse_args()


class PredictNet(nn.Module):
    def __init__(self):
        super().__init__()

        self.conv1 = nn.Sequential(
            nn.Conv2d(in_channels=1, out_channels=16, kernel_size=3, stride=1),
            nn.BatchNorm2d(num_features=16),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=2, stride=2),
            nn.Dropout2d(),
        )

        self.conv2 = nn.Sequential(
            nn.Conv2d(in_channels=16, out_channels=64, kernel_size=3, stride=1),
            nn.BatchNorm2d(num_features=64),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=2, stride=2),
            nn.Dropout2d(),
        )

        self.conv3 = nn.Sequential(
            nn.Conv2d(in_channels=64, out_channels=32, kernel_size=3, stride=1),
            nn.BatchNorm2d(num_features=32),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=2, stride=2),
            nn.Dropout2d(),
        )

        self.linear = nn.Sequential(
            nn.Linear(in_features=4 * 4 * 32, out_features=100),
            nn.ReLU(inplace=True),
            nn.Linear(in_features=100, out_features=50),
            nn.ReLU(inplace=True),
        )

        self.linear_x = nn.Linear(50, 1)
        self.linear_y = nn.Linear(50, 1)
        self.linear_sigma = nn.Linear(50, 1)

    def forward(self, x):
        BatchSize, _, _, _ = x.shape
        conv1_out = self.conv1(x)
        conv2_out = self.conv2(conv1_out)
        conv3_out = self.conv3(conv2_out)
        linear = self.linear(conv3_out.reshape(BatchSize, -1))

        return [self.linear_x(linear),
                self.linear_y(linear),
                self.linear_sigma(linear)]


class Data(Dataset):
    def __init__(self, image_size=50, data_size=10000):
        # 创建数据
        if args.cuda_id != -1:
            device = 'cuda:{}'.format(args.cuda_id)
        else:
            device = 'cpu'

        self.data_size = data_size
        self.x = torch.randint(low=0, high=image_size, size=[data_size, 1, 1],
                               device=device, dtype=torch.float)
        self.y = torch.randint(low=0, high=image_size, size=[data_size, 1, 1],
                               device=device, dtype=torch.float)
        self.sigma = torch.randint(low=30, high=100, size=[data_size, 1, 1],
                                   device=device, dtype=torch.float)

        x_grid, y_grid = torch.meshgrid(
            [torch.linspace(0, image_size, image_size, device=device),
             torch.linspace(0, image_size, image_size, device=device)],
        )

        x_grid = x_grid.unsqueeze(0).repeat_interleave(repeats=data_size, dim=0)
        y_grid = y_grid.unsqueeze(0).repeat_interleave(repeats=data_size, dim=0)

        self.gauss_map = 100 / (2 * np.pi * self.sigma) * torch.exp(
            ((x_grid - self.x) ** 2 + (y_grid - self.y) ** 2) / (2 * self.sigma ** 2)
        )

    def __getitem__(self, index):
        gauss_map = self.gauss_map[index, :, :]

        x = self.x[index, 0, 0]
        y = self.y[index, 0, 0]
        sigma = self.sigma[index, 0, 0]

        return gauss_map.unsqueeze(0), [x, y, sigma]

    def __len__(self):
        return self.data_size


def train():
    pred_net = PredictNet().cuda()

    criterion = torch.nn.MSELoss()
    optim = torch.optim.Adam(pred_net.parameters(), lr=1e-4)
    scheduler = torch.optim.lr_scheduler.StepLR(optim,1,0.9)

    dataset = Data()
    dataloader = DataLoader(dataset, batch_size=args.batch_size, shuffle=True)

    for epoch in range(args.epochs):
        print("Epoch:{epoch}, lr:{lr}".format(epoch=epoch, lr=optim.param_groups[0]['lr']))
        for gauss_map, [x, y, sigma] in dataloader:
            x_pred, y_pred, sigma_pred = pred_net(gauss_map)
            loss = criterion(x_pred, x) + criterion(y_pred, y) + criterion(sigma_pred, sigma)

            optim.zero_grad()
            loss.backward()
            optim.step()
        scheduler.step()

        # save model
        if not os.path.exists(args.model_save_path):
            os.makedirs(args.model_save_path)
        save_name = os.path.join(args.model_save_path, "c{}.pth".format(epoch))
        torch.save(pred_net.state_dict(), save_name)

        # val
        loss_val = torch.zeros([100]).cuda()
        for index, (gauss_map, [x, y, sigma]) in enumerate(dataloader):

            if index == 100:
                break

            x_pred, y_pred, sigma_pred = pred_net(gauss_map)
            loss = criterion(x_pred, x) + criterion(y_pred, y) + criterion(sigma_pred, sigma)
            loss_val[index] = loss

        print("Epoch:{epoch}, loss:{loss}".format(epoch=epoch, loss=loss_val.mean()))


def eval():
    pred_net = PredictNet().cuda()
    pred_net.load_state_dict(torch.load(args.pretrain_model_path))

    dataset = Data()
    dataloader = DataLoader(dataset, batch_size=args.batch_size, shuffle=True)
    criterion = torch.nn.MSELoss()

    loss_val = torch.zeros([100]).cuda()
    for index, (gauss_map, [x, y, sigma]) in enumerate(dataloader):
        if index == 100:
            break
        x_pred, y_pred, sigma_pred = pred_net(gauss_map)
        loss = criterion(x_pred, x) + criterion(y_pred, y) + criterion(sigma_pred, sigma)
        loss_val[index] = loss

    print("Model:{model}, loss:{loss}".format(model=args.pretrain_model_path,
                                              loss=loss_val.mean()))


def set_seed(seed=123):
    # random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)


if __name__ == '__main__':
    set_seed()
    os.environ["CUDA_VISIBLE_DEVICES"] = str(args.cuda_id)

    if args.train == "True":
        train()
    else:
        eval()

代码功能：
生成一个2D的高斯图，其中峰值位置(x,y)和方差sigma是随机的。输入这个2D高斯图，利用卷积神经网络预测这三个量。

后记

1. 尽管已经用了很久的pytorch，但是单独写这个不到200行的代码还是用了一天时间。毕竟以前只是在现成的代码上修改。
2. 数据增强部分没有包含，以后用的时候再说吧。
3. 测试的时候遇到了输出为NaN的问题，是计算结果过大或过小导致。在调整sigma的取值下限后，问题解决。