pytorch 对已有模型结构进行修改

zyxzyx_

于 2022-10-09 10:58:28 发布

阅读量972

点赞数

CC 4.0 BY-SA版权

文章标签： pytorch python 深度学习

本文链接：https://blog.csdn.net/zyxzyx_/article/details/127222439

本文介绍了如何使用PyTorch加载预训练与未预训练的VGG16模型，并展示了如何修改模型结构，包括添加新的全连接层进行迁移学习。

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

pretrained=True已经训练好的模型参数，会加载原模型

pretrained=False还未训练好的模型参数，参数为初始化的数据

vgg16分为特征层、平均池化、分类层

代码：

import torch
import torchvision
from torch import nn

vgg16_false=torchvision.models.vgg16(pretrained=False)
vgg16_true=torchvision.models.vgg16(pretrained=True)

print(vgg16_true)

#train_data=torchvision.datasets.CIFAR10('')
#在末尾增添一行线性层，降维为10（先输出原模型结构，得到1000这个原维度）
vgg16_true.add_module('add_linear',nn.Linear(1000,10))
print(vgg16_true)

加上最后一层模型如图：

#在classifier层最后面增添一层vgg16_true.classifier.add_module('add_linear',nn.Linear(1000,10))
print(vgg16_true)

图片：

修改某一层

vgg16_false.classifier[6]=nn.Linear(4096,10)
print(vgg16_false)

图片：从图1到图2