pytorch下搭建网络训练并保存模型

最近在学习pytorch，使用mnist数据集，搭建AlexNet训练并保存模型，将代码做一记录。

建立数据集的方法见pytorch建立自己的数据集（以mnist为例）

搭建网络的方法见用pytorch搭建AlexNet（微调预训练模型及手动搭建）

训练代码如下：

import torch
import os
from torchvision import transforms
import torch.optim as optim
from torch.autograd import Variable
from torch.utils.data import DataLoader
import DataProcessing as DP
import BuildModel as BM
import torch.nn as nn

if __name__ == '__main__':
    os.environ['CUDA_VISIBLE_DEVICES'] = '0,1'
    
    root_path = '/opt/Data/lixiang/ex./pytorch/Alexnet/data/'
    training_path = 'trainingset/'
    test_path = 'testset/'
    model_path = '/opt/Data/lixiang/ex./pytorch/Alexnet/model/'
    
    training_imgfile = training_path + 'trainingset_img.txt'
    training_labelfile = training_path + 'trainingset_label.txt'
    training_imgdata = training_path + 'img/'
    
    test_imgfile = test_path + 'testset_img.txt'
    test_labelfile = test_path + 'testset_label.txt'
    test_imgdata = test_path + 'img/'
    
    #parameter
    batch_size = 128
    epochs = 20
    model_type = 'pre'
    nclasses = 10
    lr = 0.01
    use_gpu = torch.cuda.is_avai