Spaces:
Sleeping
Sleeping
import torch | |
import torch.nn as nn | |
import torch.nn.functional as F | |
from climategan.blocks import InterpolateNearest2d | |
from climategan.utils import find_target_size | |
class _ASPPModule(nn.Module): | |
# https://github.com/jfzhang95/pytorch-deeplab-xception/blob/master/modeling/aspp.py | |
def __init__( | |
self, inplanes, planes, kernel_size, padding, dilation, BatchNorm, no_init | |
): | |
super().__init__() | |
self.atrous_conv = nn.Conv2d( | |
inplanes, | |
planes, | |
kernel_size=kernel_size, | |
stride=1, | |
padding=padding, | |
dilation=dilation, | |
bias=False, | |
) | |
self.bn = BatchNorm(planes) | |
self.relu = nn.ReLU() | |
if not no_init: | |
self._init_weight() | |
def forward(self, x): | |
x = self.atrous_conv(x) | |
x = self.bn(x) | |
return self.relu(x) | |
def _init_weight(self): | |
for m in self.modules(): | |
if isinstance(m, nn.Conv2d): | |
torch.nn.init.kaiming_normal_(m.weight) | |
elif isinstance(m, nn.BatchNorm2d): | |
m.weight.data.fill_(1) | |
m.bias.data.zero_() | |
class ASPP(nn.Module): | |
# https://github.com/jfzhang95/pytorch-deeplab-xception/blob/master/modeling/aspp.py | |
def __init__(self, backbone, output_stride, BatchNorm, no_init): | |
super().__init__() | |
if backbone == "mobilenet": | |
inplanes = 320 | |
else: | |
inplanes = 2048 | |
if output_stride == 16: | |
dilations = [1, 6, 12, 18] | |
elif output_stride == 8: | |
dilations = [1, 12, 24, 36] | |
else: | |
raise NotImplementedError | |
self.aspp1 = _ASPPModule( | |
inplanes, | |
256, | |
1, | |
padding=0, | |
dilation=dilations[0], | |
BatchNorm=BatchNorm, | |
no_init=no_init, | |
) | |
self.aspp2 = _ASPPModule( | |
inplanes, | |
256, | |
3, | |
padding=dilations[1], | |
dilation=dilations[1], | |
BatchNorm=BatchNorm, | |
no_init=no_init, | |
) | |
self.aspp3 = _ASPPModule( | |
inplanes, | |
256, | |
3, | |
padding=dilations[2], | |
dilation=dilations[2], | |
BatchNorm=BatchNorm, | |
no_init=no_init, | |
) | |
self.aspp4 = _ASPPModule( | |
inplanes, | |
256, | |
3, | |
padding=dilations[3], | |
dilation=dilations[3], | |
BatchNorm=BatchNorm, | |
no_init=no_init, | |
) | |
self.global_avg_pool = nn.Sequential( | |
nn.AdaptiveAvgPool2d((1, 1)), | |
nn.Conv2d(inplanes, 256, 1, stride=1, bias=False), | |
BatchNorm(256), | |
nn.ReLU(), | |
) | |
self.conv1 = nn.Conv2d(1280, 256, 1, bias=False) | |
self.bn1 = BatchNorm(256) | |
self.relu = nn.ReLU() | |
self.dropout = nn.Dropout(0.5) | |
if not no_init: | |
self._init_weight() | |
def forward(self, x): | |
x1 = self.aspp1(x) | |
x2 = self.aspp2(x) | |
x3 = self.aspp3(x) | |
x4 = self.aspp4(x) | |
x5 = self.global_avg_pool(x) | |
x5 = F.interpolate(x5, size=x4.size()[2:], mode="bilinear", align_corners=True) | |
x = torch.cat((x1, x2, x3, x4, x5), dim=1) | |
x = self.conv1(x) | |
x = self.bn1(x) | |
x = self.relu(x) | |
return self.dropout(x) | |
def _init_weight(self): | |
for m in self.modules(): | |
if isinstance(m, nn.Conv2d): | |
# n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels | |
# m.weight.data.normal_(0, math.sqrt(2. / n)) | |
torch.nn.init.kaiming_normal_(m.weight) | |
elif isinstance(m, nn.BatchNorm2d): | |
m.weight.data.fill_(1) | |
m.bias.data.zero_() | |
class DeepLabV2Decoder(nn.Module): | |
# https://github.com/jfzhang95/pytorch-deeplab-xception/blob/master/modeling/decoder.py | |
# https://github.com/jfzhang95/pytorch-deeplab-xception/blob/master/modeling/deeplab.py | |
def __init__(self, opts, no_init=False): | |
super().__init__() | |
self.aspp = ASPP("resnet", 16, nn.BatchNorm2d, no_init) | |
self.use_dada = ("d" in opts.tasks) and opts.gen.s.use_dada | |
conv_modules = [ | |
nn.Conv2d(256, 256, kernel_size=3, stride=1, padding=1, bias=False), | |
nn.BatchNorm2d(256), | |
nn.ReLU(), | |
nn.Dropout(0.5), | |
nn.Conv2d(256, 256, kernel_size=3, stride=1, padding=1, bias=False), | |
nn.BatchNorm2d(256), | |
nn.ReLU(), | |
nn.Dropout(0.1), | |
] | |
if opts.gen.s.upsample_featuremaps: | |
conv_modules = [InterpolateNearest2d(scale_factor=2)] + conv_modules | |
conv_modules += [ | |
nn.Conv2d(256, opts.gen.s.output_dim, kernel_size=1, stride=1), | |
] | |
self.conv = nn.Sequential(*conv_modules) | |
self._target_size = find_target_size(opts, "s") | |
print( | |
" - {}: setting target size to {}".format( | |
self.__class__.__name__, self._target_size | |
) | |
) | |
def set_target_size(self, size): | |
""" | |
Set final interpolation's target size | |
Args: | |
size (int, list, tuple): target size (h, w). If int, target will be (i, i) | |
""" | |
if isinstance(size, (list, tuple)): | |
self._target_size = size[:2] | |
else: | |
self._target_size = (size, size) | |
def forward(self, z, z_depth=None): | |
if self._target_size is None: | |
error = "self._target_size should be set with self.set_target_size()" | |
error += "to interpolate logits to the target seg map's size" | |
raise Exception(error) | |
if isinstance(z, (list, tuple)): | |
z = z[0] | |
if z.shape[1] != 2048: | |
raise Exception( | |
"Segmentation decoder will only work with 2048 channels for z" | |
) | |
if z_depth is not None and self.use_dada: | |
z = z * z_depth | |
y = self.aspp(z) | |
y = self.conv(y) | |
return F.interpolate(y, self._target_size, mode="bilinear", align_corners=True) | |