#!/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_pv.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)