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

sjtu_leexx

于 2018-10-18 13:56:12 发布

阅读量4.2k

点赞数 2

CC 4.0 BY-SA版权

分类专栏： pytorch python 文章标签： pytorch 神经网络

本文链接：https://blog.youkuaiyun.com/sjtuxx_lee/article/details/83145552

本文记录了使用PyTorch在MNIST数据集上搭建AlexNet网络的过程，包括数据处理、模型训练及模型保存。重点提及了数据类型的要求，如标签需为LongTensor，损失函数使用nn.CrossEntropyLoss，并强调了GPU运算时的数据迁移操作。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

最近在学习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