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

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