181 lines
7.1 KiB
Python
181 lines
7.1 KiB
Python
#!/usr/bin/env python
|
||
# -*- coding: utf-8 -*-
|
||
"""
|
||
@project:
|
||
@File : pspnet
|
||
@Author : qiqq
|
||
@create_time : 2022/11/4 22:37
|
||
"""
|
||
|
||
import torch
|
||
import torch.nn.functional as F
|
||
from torch import nn
|
||
|
||
|
||
from taihuyuan_roof.compared_experiment.pspnet.model.resnet import resnet50,resnet101
|
||
|
||
|
||
class Resnet(nn.Module):
|
||
def __init__(self, dilate_scale=8):
|
||
super(Resnet, self).__init__()
|
||
from functools import partial
|
||
model = resnet50()
|
||
# model = resnet101()
|
||
|
||
# --------------------------------------------------------------------------------------------#
|
||
# 根据下采样因子修改卷积的步长与膨胀系数
|
||
# 当downsample_factor=16的时候，我们最终获得两个特征层，shape分别是：30,30,1024和30,30,2048
|
||
# --------------------------------------------------------------------------------------------#
|
||
if dilate_scale == 8:
|
||
model.layer3.apply(partial(self._nostride_dilate, dilate=2))
|
||
model.layer4.apply(partial(self._nostride_dilate, dilate=4))
|
||
elif dilate_scale == 16:
|
||
model.layer4.apply(partial(self._nostride_dilate, dilate=2))
|
||
|
||
self.conv1 = model.conv1[0]
|
||
self.bn1 = model.conv1[1]
|
||
self.relu1 = model.conv1[2]
|
||
self.conv2 = model.conv1[3]
|
||
self.bn2 = model.conv1[4]
|
||
self.relu2 = model.conv1[5]
|
||
self.conv3 = model.conv1[6]
|
||
self.bn3 = model.bn1
|
||
self.relu3 = model.relu
|
||
self.maxpool = model.maxpool
|
||
self.layer1 = model.layer1
|
||
self.layer2 = model.layer2
|
||
self.layer3 = model.layer3
|
||
self.layer4 = model.layer4
|
||
|
||
def _nostride_dilate(self, m, dilate):
|
||
classname = m.__class__.__name__
|
||
if classname.find('Conv') != -1:
|
||
if m.stride == (2, 2):
|
||
m.stride = (1, 1)
|
||
if m.kernel_size == (3, 3):
|
||
m.dilation = (dilate // 2, dilate // 2)
|
||
m.padding = (dilate // 2, dilate // 2)
|
||
else:
|
||
if m.kernel_size == (3, 3):
|
||
m.dilation = (dilate, dilate)
|
||
m.padding = (dilate, dilate)
|
||
|
||
def forward(self, x):
|
||
x = self.relu1(self.bn1(self.conv1(x)))
|
||
x = self.relu2(self.bn2(self.conv2(x)))
|
||
x = self.relu3(self.bn3(self.conv3(x)))
|
||
x = self.maxpool(x)
|
||
|
||
x = self.layer1(x)
|
||
x = self.layer2(x)
|
||
x_aux = self.layer3(x)
|
||
x = self.layer4(x_aux)
|
||
return x_aux, x
|
||
|
||
|
||
|
||
|
||
class _PSPModule(nn.Module):
|
||
def __init__(self, in_channels, pool_sizes, norm_layer):
|
||
super(_PSPModule, self).__init__()
|
||
out_channels = in_channels // len(pool_sizes)
|
||
# -----------------------------------------------------#
|
||
# 分区域进行平均池化
|
||
# 30, 30, 320 + 30, 30, 80 + 30, 30, 80 + 30, 30, 80 + 30, 30, 80 = 30, 30, 640
|
||
# -----------------------------------------------------#
|
||
self.stages = nn.ModuleList(
|
||
[self._make_stages(in_channels, out_channels, pool_size, norm_layer) for pool_size in pool_sizes])
|
||
|
||
# 30, 30, 640 -> 30, 30, 80
|
||
self.bottleneck = nn.Sequential(
|
||
nn.Conv2d(in_channels + (out_channels * len(pool_sizes)), out_channels, kernel_size=3, padding=1,
|
||
bias=False),
|
||
norm_layer(out_channels),
|
||
nn.ReLU(inplace=True),
|
||
nn.Dropout2d(0.1)
|
||
)
|
||
|
||
def _make_stages(self, in_channels, out_channels, bin_sz, norm_layer):
|
||
prior = nn.AdaptiveAvgPool2d(output_size=bin_sz)
|
||
conv = nn.Conv2d(in_channels, out_channels, kernel_size=1, bias=False)
|
||
bn = norm_layer(out_channels)
|
||
relu = nn.ReLU(inplace=True)
|
||
return nn.Sequential(prior, conv, bn, relu)
|
||
|
||
def forward(self, features):#以512*512 降采样16倍 到这里是32*32*2048
|
||
h, w = features.size()[2], features.size()[3]
|
||
pyramids = [features] #这个地方这样写就成了5个分支了 我觉得应该还是[]
|
||
pyramids.extend(
|
||
[F.interpolate(stage(features), size=(h, w), mode='bilinear', align_corners=True) for stage in self.stages])
|
||
output = self.bottleneck(torch.cat(pyramids, dim=1)) #4096降维到 512
|
||
return output
|
||
|
||
|
||
class PSPNet(nn.Module):
|
||
def __init__(self, num_classes, downsample_factor, backbone="resnet50", aux_branch=False):
|
||
super(PSPNet, self).__init__()
|
||
norm_layer = nn.BatchNorm2d
|
||
if backbone == "resnet50":
|
||
self.backbone = Resnet(downsample_factor)
|
||
aux_channel = 1024
|
||
out_channel = 2048
|
||
|
||
|
||
# --------------------------------------------------------------#
|
||
# PSP模块，分区域进行池化
|
||
# 分别分割成1x1的区域，2x2的区域，3x3的区域，6x6的区域
|
||
# 30,30,320 -> 30,30,80 -> 30,30,21
|
||
# --------------------------------------------------------------#
|
||
self.master_branch = nn.Sequential(
|
||
_PSPModule(out_channel, pool_sizes=[1, 2, 3, 6], norm_layer=norm_layer),
|
||
nn.Conv2d(out_channel // 4, num_classes, kernel_size=1)
|
||
)
|
||
|
||
self.aux_branch = aux_branch
|
||
|
||
if self.aux_branch:
|
||
# ---------------------------------------------------#
|
||
# 利用特征获得预测结果
|
||
# 30, 30, 96 -> 30, 30, 40 -> 30, 30, 21
|
||
# ---------------------------------------------------#
|
||
self.auxiliary_branch = nn.Sequential(
|
||
nn.Conv2d(aux_channel, out_channel // 8, kernel_size=3, padding=1, bias=False),
|
||
norm_layer(out_channel // 8),
|
||
nn.ReLU(inplace=True),
|
||
nn.Dropout2d(0.1),
|
||
nn.Conv2d(out_channel // 8, num_classes, kernel_size=1)
|
||
)
|
||
|
||
self.initialize_weights(self.master_branch)
|
||
|
||
def forward(self, x):
|
||
input_size = (x.size()[2], x.size()[3])
|
||
x_aux, x = self.backbone(x)
|
||
output = self.master_branch(x)
|
||
output = F.interpolate(output, size=input_size, mode='bilinear', align_corners=True)
|
||
if self.aux_branch:
|
||
output_aux = self.auxiliary_branch(x_aux)
|
||
output_aux = F.interpolate(output_aux, size=input_size, mode='bilinear', align_corners=True)
|
||
return output_aux, output
|
||
else:
|
||
return output
|
||
|
||
def initialize_weights(self, *models):
|
||
for model in models:
|
||
for m in model.modules():
|
||
if isinstance(m, nn.Conv2d):
|
||
nn.init.kaiming_normal_(m.weight.data, nonlinearity='relu')
|
||
elif isinstance(m, nn.BatchNorm2d):
|
||
m.weight.data.fill_(1.)
|
||
m.bias.data.fill_(1e-4)
|
||
elif isinstance(m, nn.Linear):
|
||
m.weight.data.normal_(0.0, 0.0001)
|
||
m.bias.data.zero_()
|
||
|
||
|
||
if __name__ == '__main__':
|
||
inpuu =torch.rand(2,3,512,512)
|
||
model =PSPNet(num_classes=3,downsample_factor=8,)
|
||
# model.eval()
|
||
out =model(inpuu)
|
||
print(out.shape) |