| import torch
|
| import torch.nn as nn
|
| import torch.nn.functional as functional
|
|
|
| from models._util import try_index
|
| from .bn import ABN
|
|
|
|
|
| class DeeplabV3(nn.Module):
|
| def __init__(self,
|
| in_channels,
|
| out_channels,
|
| hidden_channels=256,
|
| dilations=(12, 24, 36),
|
| norm_act=ABN,
|
| pooling_size=None):
|
| super(DeeplabV3, self).__init__()
|
| self.pooling_size = pooling_size
|
|
|
| self.map_convs = nn.ModuleList([
|
| nn.Conv2d(in_channels, hidden_channels, 1, bias=False),
|
| nn.Conv2d(in_channels, hidden_channels, 3, bias=False, dilation=dilations[0], padding=dilations[0]),
|
| nn.Conv2d(in_channels, hidden_channels, 3, bias=False, dilation=dilations[1], padding=dilations[1]),
|
| nn.Conv2d(in_channels, hidden_channels, 3, bias=False, dilation=dilations[2], padding=dilations[2])
|
| ])
|
| self.map_bn = norm_act(hidden_channels * 4)
|
|
|
| self.global_pooling_conv = nn.Conv2d(in_channels, hidden_channels, 1, bias=False)
|
| self.global_pooling_bn = norm_act(hidden_channels)
|
|
|
| self.red_conv = nn.Conv2d(hidden_channels * 4, out_channels, 1, bias=False)
|
| self.pool_red_conv = nn.Conv2d(hidden_channels, out_channels, 1, bias=False)
|
| self.red_bn = norm_act(out_channels)
|
|
|
| self.reset_parameters(self.map_bn.activation, self.map_bn.slope)
|
|
|
| def reset_parameters(self, activation, slope):
|
| gain = nn.init.calculate_gain(activation, slope)
|
| for m in self.modules():
|
| if isinstance(m, nn.Conv2d):
|
| nn.init.xavier_normal_(m.weight.data, gain)
|
| if hasattr(m, "bias") and m.bias is not None:
|
| nn.init.constant_(m.bias, 0)
|
| elif isinstance(m, ABN):
|
| if hasattr(m, "weight") and m.weight is not None:
|
| nn.init.constant_(m.weight, 1)
|
| if hasattr(m, "bias") and m.bias is not None:
|
| nn.init.constant_(m.bias, 0)
|
|
|
| def forward(self, x):
|
|
|
| out = torch.cat([m(x) for m in self.map_convs], dim=1)
|
| out = self.map_bn(out)
|
| out = self.red_conv(out)
|
|
|
|
|
| pool = self._global_pooling(x)
|
| pool = self.global_pooling_conv(pool)
|
| pool = self.global_pooling_bn(pool)
|
| pool = self.pool_red_conv(pool)
|
| if self.training or self.pooling_size is None:
|
| pool = pool.repeat(1, 1, x.size(2), x.size(3))
|
|
|
| out += pool
|
| out = self.red_bn(out)
|
| return out
|
|
|
| def _global_pooling(self, x):
|
| if self.training or self.pooling_size is None:
|
| pool = x.view(x.size(0), x.size(1), -1).mean(dim=-1)
|
| pool = pool.view(x.size(0), x.size(1), 1, 1)
|
| else:
|
| pooling_size = (min(try_index(self.pooling_size, 0), x.shape[2]),
|
| min(try_index(self.pooling_size, 1), x.shape[3]))
|
| padding = (
|
| (pooling_size[1] - 1) // 2,
|
| (pooling_size[1] - 1) // 2 if pooling_size[1] % 2 == 1 else (pooling_size[1] - 1) // 2 + 1,
|
| (pooling_size[0] - 1) // 2,
|
| (pooling_size[0] - 1) // 2 if pooling_size[0] % 2 == 1 else (pooling_size[0] - 1) // 2 + 1
|
| )
|
|
|
| pool = functional.avg_pool2d(x, pooling_size, stride=1)
|
| pool = functional.pad(pool, pad=padding, mode="replicate")
|
| return pool
|
|
|
