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

最新推荐文章于 2025-03-18 19:36:33 发布

原创

最新推荐文章于 2025-03-18 19:36:33 发布 · 4.3k 阅读

15 ·

CC 4.0 BY-SA版权

文章标签：

#pytorch #神经网络

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

最近在学习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_available()
    
    transformations = transforms

最低0.47元/天解锁文章