【pytorch修改预训练pt文件】

---

0. 前言

---

1. nn.module

1.1 使用nn.module的model.children()的函数，重新定义自己model的层

那么对于网络的修改怎样可以快速的去除model本身的网络层呢？
一个继承nn.module的model它包含一个叫做children()的函数，这个函数可以用来提取出model每一层的网络结构，在此基础上进行修改即可，修改方法如下(去除后两层)：

1 resnet_layer = nn.Sequential(*list(model.children())[:-2])

那么，接下来就可以构建我们的网络了：：

class Net(nn.Module):
    def __init__(self , model):	#此处的model参数是已经加载了预训练参数的模型，方便继承预训练成果
        super(Net, self).__init__()
        #取掉model的后两层
        self.resnet_layer = nn.Sequential(*list(model.children())[:-2])
        
        self.transion_layer = nn.ConvTranspose2d(2048, 2048, kernel_size=14, stride=3)
        
        
    def forward(self, x):
        x = self.resnet_layer(x)
        x = self.transion_layer(x)
        return x

2. Yolov5

2.1 nn.Module def函数

def _conv_bn(input_channel, output_channel, kernel_size=1, stride=1, padding=1, groups=1):
    res = nn.Sequential()
    res.add_module('conv',nn.Conv2d(in_channels=input_channel, out_channels=output_channel, kernel_size=kernel_size,padding=padding, padding_mode='zeros', stride=stride, groups=groups, bias=False))
    res.add_module('bn', nn.BatchNorm2d(output_channel))
    return res
'''我之前是把_conv_bn()函数写进了Con_C3类中，
但是发现self.brb_3x3 =self._conv_bn(c1, c2, kernel_size=k, padding=autopad(k, p),groups=g)
self.brb_1x1 = self._conv_bn(c1, c2, kernel_size=1, padding=0, groups=g)
这样会出现_conv_bn() got multiple values for argument 'kernel_size'错误，因为self中的函数就只有一个，所以必须要把_conv_bn放在Conv_C3外面

'''
class Conv_C3(nn.Module):
    # Standard convolution
    def __init__(self, c1, c2, k=1, s=1,deploy=False, p=None, g=1, act=True):  # ch_in, ch_out, kernel, stride, padding, groups
        super().__init__()
        self.conv = nn.Conv2d(c1, c2, k, s, autopad(k, p), groups=g, bias=False)
        self.bn = nn.BatchNorm2d(c2)
        self.act = nn.SiLU() if act is True else (act if isinstance(act, nn.Module) else nn.Identity())
        self.deploy = deploy
        self.input_channel = c1
        self.k =  k
        # class Conv_non_deep(nn.Module):
        if (not self.deploy):#not  deploy=detect
            self.brb_3x3 = _conv_bn(c1, c2, kernel_size=k, padding=autopad(k, p),groups=g)
            self.brb_1x1 = _conv_bn(c1, c2, kernel_size=1, padding=0, groups=g)
            self.brb_identity = nn.BatchNorm2d(c1) if c1 == c2 else None   #这个在接下来的C3处会造成self.cv3无法重参数化
        else:  #detect

#experimental.py
for m in model.modules():
	if type(m) in [nn.Hardwish,nn.LeakyRelu,...]:
		xxx
	elif type(m) is Conv:
		m._non_persistent_buffers_set = set()
		#这个是直接可以通过m.操作来进行修改吗？

2.2 torch.save

torch.save(model,'save.pt')#保存整个模型
torch.save(model.state_dict(), 'save.pt')#只保存训练好的权重

3. nn.Sequential

3.1 Sequential 序贯模型

import torch.nn as nn
model = nn.Sequential(
                  nn.Conv2d(1,20,5),
                  nn.ReLU(),
                  nn.Conv2d(20,64,5),
                  nn.ReLU()
                )
 
print(model)
print(model[2]) # 通过索引获取第几个层
'''运行结果为：
Sequential(
  (0): Conv2d(1, 20, kernel_size=(5, 5), stride=(1, 1))
  (1): ReLU()
  (2): Conv2d(20, 64, kernel_size=(5, 5), stride=(1, 1))
  (3): ReLU()
)
Conv2d(20, 64, kernel_size=(5, 5), stride=(1, 1))
'''

3.2 给每一个层添加名称

import torch.nn as nn
from collections import OrderedDict
model = nn.Sequential(OrderedDict([
                  ('conv1', nn.Conv2d(1,20,5)),
                  ('relu1', nn.ReLU()),
                  ('conv2', nn.Conv2d(20,64,5)),
                  ('relu2', nn.ReLU())
                ]))
 
print(model)
print(model[2]) # 通过索引获取第几个层
'''运行结果为：
Sequential(
  (conv1): Conv2d(1, 20, kernel_size=(5, 5), stride=(1, 1))
  (relu1): ReLU()
  (conv2): Conv2d(20, 64, kernel_size=(5, 5), stride=(1, 1))
  (relu2): ReLU()
)
Conv2d(20, 64, kernel_size=(5, 5), stride=(1, 1))
'''

3.3 nn.Sequential().add_module

nn.Sequential().add_module(‘conv1’,nn.Conv2d(1,20,5))

import torch.nn as nn
from collections import OrderedDict
model = nn.Sequential()
model.add_module("conv1",nn.Conv2d(1,20,5))
model.add_module('relu1', nn.ReLU())
model.add_module('conv2', nn.Conv2d(20,64,5))
model.add_module('relu2', nn.ReLU())
 
print(model)
print(model[2]) # 通过索引获取第几个层