Pytorch训练神经网络样例代码 -- 单文件版

这篇博客使用LeNet网络结构针对数据集Cifar10训练深度神经网络分类器。主要是作为以后编码神经网络代码的一种编程习惯。尤其注意其中封装的根据gpu_ids自动使用相关GPU进行训练的代码模块。

  【 train.py 】

import torch as t
import torchvision as tv
import torchvision.transforms as transforms
import torch.nn as nn
import torch.nn.functional as F


########## 超参数设置 ##########
epochs = 200
learning_rate = 0.001
batch_size = 256
gpu_ids = [0, 1, 2]


########## 一、数据加载与预处理 ##########
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])
trainset = tv.datasets.CIFAR10(root='./data/', train=True, download=True, transform=transform)
testset = tv.datasets.CIFAR10(root='./data/', train=False, download=True, transform=transform)
trainloader = t.utils.data.DataLoader(trainset, batch_size=batch_size, shuffle=True, num_workers=4)
testloader = t.utils.data.DataLoader(testset, batch_size=batch_size, shuffle=False, num_workers=4)


########## 二、定义神经网络 ##########
class LeNet(nn.Module):
    def __init__(self):
        super(LeNet, self).__init__()
        # (in_channels, out_channels, k