pytorch网络的增删改

本文介绍对加载的网络的层进行增删改, 以alexnet网络为例进行介绍。

1. 加载网络

import torchvision.models as models  

alexnet =models.alexnet(weights=models.AlexNet_Weights.DEFAULT)
print(alexnet)

2. 删除网络

在做迁移学习的时候，我们通常是在分类网络的基础上进行修改的。一般会把网络最后的几层删除掉，主要是全局平均池化层、全连接层。只留前面的网络部分作为特征提取器，再次基础上进行其他的任务。

2.1 删除网络任意层

• 将alexnet的classifier这一部分全删除掉

del  alexnet.classifer
print(alexnet)

删除classifer模块后，打印结果如下：

可以看到只剩下features和avgpool这两个模块了。刚才的classifier就已经被我们删除掉了。

• 删除classifier模块中的某一层

如果不想把classifier这一模块整体删除，只想删除classifier中比如第6个层

# del alexnet.classifier 
del alexnet.classifier[6]
print(alexnet)

可以看到classifier中第6层就已经被删除掉了。

2.2 删除网络的最后多层

如果想把网络的连续几层给删除掉，比如classifier中最后的几层删除掉

#------------------删除网络的最后多层--------------------------#
alexnet.classifier = alexnet.classifier[:-2]
print(alexnet)
#-------------------------------------------------------------#

打印信息如下：

• 可以看出classifier看出最后2层(5,6)被删除掉了

可以使用切片的方式，保留不需要被删除的层重新赋给classifier模块，没有保留的就被删除了。

3. 修改网络的某一层

• 没有修改之前alexnet.classifier的第6层是个全连接层，输入通道为4096, 输出通道为1000。
  
• 假设此时，我们想最后一层全连接层的输出，改为1024。此时，你只需要重新定义这层全连接层。

#-----------------修改网络的某一层-----------------------------#
alexnet.classifier[6] = nn.Linear(in_features=4096,out_features=1024)
print(alexnet)
#-------------------------------------------------------------#

打印后，可以看到最后一层的输出由原来的4096改为了1024

4. 在网络中添加某一层

4.1 每次添加一层

假设我们想在网络最后输出中，再添加两层，分别为ReLU和nn.Linear层

#-----------------修改网络的某一层-----------------------------#
alexnet.classifier[6] = nn.Linear(in_features=4096,out_features=1024)
# print(alexnet)
#-------------------------------------------------------------#


#-------------网络添加层，每次添加一层--------------------------#
alexnet.classifier.add_module('7',nn.ReLU(inplace=True))
alexnet.classifier.add_module('8',nn.Linear(in_features=1024,out_features=20))
print(alexnet)
#-------------------------------------------------------------#

• 利用add_module来添加层，第一个参数为层名称，第二个参数为定义layer的内容
• 我们在alexnet.classifier这个block中进行添加的，添加后打印网络结构如下：
  
• 可以看到成功的添加了最后2层。

4.2 一次添加多层

如果觉得一层层的添加层比较麻烦，比如我们可以一次性添加一个大的模块new_block。

block = nn.Sequential(
    nn.ReLU(inplace=True),
    nn.Linear(in_features=1024,out_features=20)
)

alexnet.add_module('new_block',block)
print(alexnet)

• 可以看到在alexnet网络中新增了new_block,该block中包括2层，分别是ReLU层以及Linear层。

以上就是对Pytorch网络增删改的方法，完整的代码如下：

import torchvision.models as models  
import torch.nn   as nn 

alexnet =models.alexnet(weights=models.AlexNet_Weights.DEFAULT)
# print(alexnet)

#1.-----------------删除网络的最后一层-------------------------#
# del alexnet.classifier 
# del alexnet.classifier[6]
# print(alexnet)
#-------------------------------------------------------------#

#------------------删除网络的最后多层--------------------------#
# alexnet.classifier = alexnet.classifier[:-2]
# print(alexnet)
#-------------------------------------------------------------#

#-----------------修改网络的某一层-----------------------------#
alexnet.classifier[6] = nn.Linear(in_features=4096,out_features=1024)
# print(alexnet)
#-------------------------------------------------------------#


#-------------网络添加层，每次添加一层--------------------------#
# alexnet.classifier.add_module('7',nn.ReLU(inplace=True))
# alexnet.classifier.add_module('8',nn.Linear(in_features=1024,out_features=20))
# print(alexnet)
#-------------------------------------------------------------#

#-----------------------网络添加层，一次性添加多层--------------#
block = nn.Sequential(
    nn.ReLU(inplace=True),
    nn.Linear(in_features=1024,out_features=20)
)

alexnet.add_module('new_block',block)
print(alexnet)
#-------------------------------------------------------------#