pytorch修改预训练模型（输入通道数）

Pytorch修改预训练模型

  torchvision模块带有很多预训练模型，具体支持的模型列表可以参看官方文档。
  在语义分割中，预训练模型一般是3通道的，但是在实际情况中经常会有输入通道数量不止3个通道，要修改预训练模型的通道数,要么重写模型，要么就复用预训练模型，这里推荐用官方的预训练模型，代码简洁且不容易出错。
  步入正题，要调整预训练模型需要两个步骤，首先如下加载预训练模型，并打印模型第一层，然后修改第一层结构的输入通道数即可。

import torchvision.models as models
backbone = models.resnet101(pretrained=False)
print(backbone.conv1)

  这里打印得到第一层的结构。

Conv2d(3, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False)

  于是这里可以看到输入模型的通道数为3，只要修改这个3就可以了，若期望的输入是4通道，则如下操作即可：

backbone.conv1= nn.Conv2d(4, 64, kernel_size=7, stride=2, padding=3,bias=False)

  至此已经修改完毕，下面是deeplabv3+的backbone的调整实例。
  这里有个小技巧，在模型调整完毕后需要先做单元测试，通过随机生成一个tensor，输入模型看能不能跑出结果（代码最后的测试模块）。

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

class ResNet_101(nn.Module):
	"""加载预训练模型resnet101"""
	 def __init__(self, in_channels = 13, conv1_out = 64):
	   super(ResNet_101,self).__init__()
	   backbone = models.resnet101(pretrained=False)
	   print("load pretrained ResNet101 model")
	   # self.conv1 = backbone.conv1
	   self.conv1 = nn.Conv2d(in_channels, 64, kernel_size=7, stride=2, padding=3,bias=False)
	   self.bn1 = backbone.bn1
	   self.relu = nn.ReLU(inplace=True)
	   self.maxpool = backbone.maxpool
	   self.layer1 = backbone.layer1
	   self.layer2 = backbone.layer2
	   self.layer3 = backbone.layer3
	 
	 def forward(self,x):
	   x = self.relu(self.bn1(self.conv1(x)))
	   x = self.maxpool(x)
	   x = self.layer1(x)
	   x = self.layer2(x)
	   x = self.layer3(x)
	   
	   return x
    
class Deeplabv3Resnet101(nn.Module):
	"""语义分割中以resnet101为backbone的deeplabv3+模型"""
    def __init__(self,nc=2,input_channel=3):
        super(deeplabv3plus_resnet101,self).__init__()
        self.nc = nc
        self.backbone = ResNet_101(input_channel)
        self.assp = ASSP(in_channels=1024)
        self.out1 = nn.Sequential(nn.Conv2d(in_channels=256,out_channels=256,kernel_size=1,stride=1),nn.ReLU())
        self.dropout1 = nn.Dropout(0.5)
        self.up4 = nn.Upsample(scale_factor=4)
        self.up2 = nn.Upsample(scale_factor=2)

        self.conv1x1 = nn.Sequential(nn.Conv2d(1024,256,1,bias=False),nn.ReLU())
        self.conv3x3 = nn.Sequential(nn.Conv2d(512,self.nc,1),nn.ReLU())
        self.dec_conv = nn.Sequential(nn.Conv2d(256,256,3,padding=1),nn.ReLU())
        
    def forward(self,x):
        x = self.backbone(x)
        out1 = self.assp(x)
        out1 = self.out1(out1)
        out1 = self.dropout1(out1)
        out1 = self.up4(out1)
        # print(out1.shape)

        dec = self.conv1x1(x)
        dec = self.dec_conv(dec)
        dec = self.up4(dec)
        concat = torch.cat((out1,dec),dim=1)
        out = self.conv3x3(concat)
        out = self.up4(out)
        return out

class ASSP(nn.Module):
    def __init__(self,in_channels,out_channels = 256):
        super(ASSP,self).__init__()
        self.relu = nn.ReLU(inplace=True)
        self.conv1 = nn.Conv2d(in_channels,out_channels,1,padding = 0,dilation=1,bias=False)
        self.bn1 = nn.BatchNorm2d(out_channels)
        self.conv2 = nn.Conv2d(in_channels = in_channels,out_channels = out_channels,kernel_size = 3,stride=1,padding = 6,dilation = 6,bias=False)
        self.bn2 = nn.BatchNorm2d(out_channels)
        self.conv3 = nn.Conv2d(in_channels = in_channels,out_channels = out_channels,kernel_size = 3,stride=1,padding = 12,dilation = 12,bias=False)
        self.bn3 = nn.BatchNorm2d(out_channels)
        self.conv4 = nn.Conv2d(in_channels = in_channels,out_channels = out_channels,kernel_size = 3,stride=1,padding = 18,dilation = 18,bias=False)
        self.bn4 = nn.BatchNorm2d(out_channels)
        self.conv5 = nn.Conv2d(in_channels = in_channels,out_channels = out_channels,kernel_size = 1,stride=1,padding = 0,dilation=1,bias=False)
        self.bn5 = nn.BatchNorm2d(out_channels)
        self.convf = nn.Conv2d(in_channels = out_channels * 5,out_channels = out_channels,kernel_size = 1,stride=1,padding = 0,dilation=1,bias=False)
        self.bnf = nn.BatchNorm2d(out_channels)
        self.adapool = nn.AdaptiveAvgPool2d(1)
    def forward(self,x):
        x1 = self.conv1(x)
        x1 = self.bn1(x1)
        x1 = self.relu(x1)
        x2 = self.conv2(x)
        x2 = self.bn2(x2)
        x2 = self.relu(x2)
        x3 = self.conv3(x)
        x3 = self.bn3(x3)
        x3 = self.relu(x3)
        x4 = self.conv4(x)
        x4 = self.bn4(x4)
        x4 = self.relu(x4)
        x5 = self.adapool(x)
        x5 = self.conv5(x5)
        x5 = self.bn5(x5)
        x5 = self.relu(x5)
        x5 = F.interpolate(x5, size = tuple(x4.shape[-2:]), mode='bilinear',align_corners=True)
    
        x = torch.cat((x1,x2,x3,x4,x5), dim = 1) #channels first
        x = self.convf(x)
        x = self.bnf(x)
        x = self.relu(x)
        return x
        
if __name__ == "__main__":
	input_tensor = torch.rand(4, 4, 128, 128)  # batch_size,input_channel,input_h,input_w
	model = Deeplabv3Resnet101()
	out = model(input_tensor)
	print(out)