模型的改进、保存、加载

G鲲鹏展翅Y

于 2024-03-11 09:53:27 发布

阅读量440

点赞数 7

文章标签： python 深度学习 pytorch

本文链接：https://blog.csdn.net/qq_44665922/article/details/136615582

版权

import torch
import torchvision
from torch import nn
import warnings

warnings.filterwarnings("ignore")
vgg16_false = torchvision.models.vgg16(weights=False)
vgg16_true = torchvision.models.vgg16(weights='VGG16_Weights.DEFAULT')   #加载默认权重的VGG16
print(vgg16_true)

vgg16_true.add_module('add_linear',nn.Linear(1000,10))   #在最后加上一个线性层
print(vgg16_true)

torch.save(vgg16_true, "model1")   #参数与模型结构都保存
torch.save(vgg16_true.state_dict(), "model2")   #只保留参数，适合大模型的保存
vgg16_load1 = torch.load("model1")
vgg16_load2 = torch.load("model2")
print(vgg16_load1,vgg16_load2)