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)