#import encoding.nn as nn #import encoding.functions as F import torch.nn as nn from torch.nn import functional as F import math import torch.utils.model_zoo as model_zoo import torch import numpy as np from torch.autograd import Variable affine_par = True import functools import sys, os from libs import InPlaceABN, InPlaceABNSync from cc_attention import CrissCrossAttention BatchNorm2d = functools.partial(InPlaceABNSync, activation='none') def outS(i): i = int(i) i = (i+1)/2 i = int(np.ceil((i+1)/2.0)) i = (i+1)/2 return i def conv3x3(in_planes, out_planes, stride=1): "3x3 convolution with padding" return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride, padding=1, bias=False) class Bottleneck(nn.Module): expansion = 4 def __init__(self, inplanes, planes, stride=1, dilation=1, downsample=None, fist_dilation=1, multi_grid=1): super(Bottleneck, self).__init__() self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False) self.bn1 = BatchNorm2d(planes) self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=stride, padding=dilation*multi_grid, dilation=dilation*multi_grid, bias=False) self.bn2 = BatchNorm2d(planes) self.conv3 = nn.Conv2d(planes, planes * 4, kernel_size=1, bias=False) self.bn3 = BatchNorm2d(planes * 4) self.relu = nn.ReLU(inplace=False) self.relu_inplace = nn.ReLU(inplace=True) self.downsample = downsample self.dilation = dilation self.stride = stride def forward(self, x): residual = x out = self.conv1(x) out = self.bn1(out) out = self.relu(out) out = self.conv2(out) out = self.bn2(out) out = self.relu(out) out = self.conv3(out) out = self.bn3(out) if self.downsample is not None: residual = self.downsample(x) out = out + residual out = self.relu_inplace(out) return out class PSPModule(nn.Module): """ Reference: Zhao, Hengshuang, et al. *"Pyramid scene parsing network."* """ def __init__(self, features, out_features=512, sizes=(1, 2, 3, 6)): super(PSPModule, self).__init__() self.stages = [] self.stages = nn.ModuleList([self._make_stage(features, out_features, size) for size in sizes]) self.bottleneck = nn.Sequential( nn.Conv2d(features+len(sizes)*out_features, out_features, kernel_size=3, padding=1, dilation=1, bias=False), InPlaceABNSync(out_features), nn.Dropout2d(0.1) ) def _make_stage(self, features, out_features, size): prior = nn.AdaptiveAvgPool2d(output_size=(size, size)) conv = nn.Conv2d(features, out_features, kernel_size=1, bias=False) bn = InPlaceABNSync(out_features) return nn.Sequential(prior, conv, bn) def forward(self, feats): h, w = feats.size(2), feats.size(3) priors = [F.upsample(input=stage(feats), size=(h, w), mode='bilinear', align_corners=True) for stage in self.stages] + [feats] bottle = self.bottleneck(torch.cat(priors, 1)) return bottle class RCCAModule(nn.Module): def __init__(self, in_channels, out_channels, num_classes): super(RCCAModule, self).__init__() inter_channels = in_channels // 4 self.conva = nn.Sequential(nn.Conv2d(in_channels, inter_channels, 3, padding=1, bias=False), InPlaceABNSync(inter_channels)) self.cca = CrissCrossAttention(inter_channels) self.convb = nn.Sequential(nn.Conv2d(inter_channels, inter_channels, 3, padding=1, bias=False), InPlaceABNSync(inter_channels)) self.bottleneck = nn.Sequential( nn.Conv2d(in_channels+inter_channels, out_channels, kernel_size=3, padding=1, dilation=1, bias=False), InPlaceABNSync(out_channels), nn.Dropout2d(0.1), nn.Conv2d(512, num_classes, kernel_size=1, stride=1, padding=0, bias=True) ) def forward(self, x, recurrence=1): output = self.conva(x) for i in range(recurrence): output = self.cca(output) output = self.convb(output) output = self.bottleneck(torch.cat([x, output], 1)) return output class ResNet(nn.Module): def __init__(self, block, layers, num_classes): self.inplanes = 128 super(ResNet, self).__init__() self.conv1 = conv3x3(3, 64, stride=2) self.bn1 = BatchNorm2d(64) self.relu1 = nn.ReLU(inplace=False) self.conv2 = conv3x3(64, 64) self.bn2 = BatchNorm2d(64) self.relu2 = nn.ReLU(inplace=False) self.conv3 = conv3x3(64, 128) self.bn3 = BatchNorm2d(128) self.relu3 = nn.ReLU(inplace=False) self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1) self.relu = nn.ReLU(inplace=False) self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1, ceil_mode=True) # change self.layer1 = self._make_layer(block, 64, layers[0]) self.layer2 = self._make_layer(block, 128, layers[1], stride=2) self.layer3 = self._make_layer(block, 256, layers[2], stride=1, dilation=2) self.layer4 = self._make_layer(block, 512, layers[3], stride=1, dilation=4, multi_grid=(1,1,1)) #self.layer5 = PSPModule(2048, 512) self.head = RCCAModule(2048, 512, num_classes) self.dsn = nn.Sequential( nn.Conv2d(1024, 512, kernel_size=3, stride=1, padding=1), InPlaceABNSync(512), nn.Dropout2d(0.1), nn.Conv2d(512, num_classes, kernel_size=1, stride=1, padding=0, bias=True) ) def get_learnable_parameters(self, freeze_layers=[True,True,True,True,False,False,False]): lr_parameters = [] if not freeze_layers[0]: for i in [self.conv1, self.bn1, self.conv2, self.bn2, self.conv3, self.bn3]: params = i.named_parameters() for name, p in params: print(name) lr_parameters.append(p) layers = [self.layer1, self.layer2, self.layer3, self.layer4, self.layer5, self.layer6] for freeze, layer in zip(freeze_layers[1:], layers): if not freeze: params = layer.named_parameters() for name, p in params: print(name) lr_parameters.append(p) return lr_parameters def _make_layer(self, block, planes, blocks, stride=1, dilation=1, multi_grid=1): downsample = None if stride != 1 or self.inplanes != planes * block.expansion: downsample = nn.Sequential( nn.Conv2d(self.inplanes, planes * block.expansion, kernel_size=1, stride=stride, bias=False), BatchNorm2d(planes * block.expansion,affine = affine_par)) layers = [] generate_multi_grid = lambda index, grids: grids[index%len(grids)] if isinstance(grids, tuple) else 1 layers.append(block(self.inplanes, planes, stride,dilation=dilation, downsample=downsample, multi_grid=generate_multi_grid(0, multi_grid))) self.inplanes = planes * block.expansion for i in range(1, blocks): layers.append(block(self.inplanes, planes, dilation=dilation, multi_grid=generate_multi_grid(i, multi_grid))) return nn.Sequential(*layers) def forward(self, x, recurrence=1): 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 = self.layer3(x) x_dsn = self.dsn(x) x = self.layer4(x) x = self.head(x, recurrence) return [x, x_dsn] def Res_Deeplab(num_classes=21): model = ResNet(Bottleneck,[3, 4, 23, 3], num_classes) return model