climategan/deeplab/deeplab_v3.py

"""
https://github.com/jfzhang95/pytorch-deeplab-xception/blob/master/modeling/backbone/resnet.py
"""
from pathlib import Path

import torch
import torch.nn as nn
import torch.nn.functional as F
from climategan.deeplab.mobilenet_v3 import SeparableConv2d
from climategan.utils import find_target_size


class _DeepLabHead(nn.Module):
    def __init__(
        self, nclass, c1_channels=256, c4_channels=2048, norm_layer=nn.BatchNorm2d
    ):
        super().__init__()
        last_channels = c4_channels
        # self.c1_block = _ConvBNReLU(c1_channels, 48, 1, norm_layer=norm_layer)
        # last_channels += 48
        self.block = nn.Sequential(
            SeparableConv2d(
                last_channels, 256, 3, norm_layer=norm_layer, relu_first=False
            ),
            SeparableConv2d(256, 256, 3, norm_layer=norm_layer, relu_first=False),
            nn.Conv2d(256, nclass, 1),
        )

    def forward(self, x, c1=None):
        return self.block(x)


class ConvBNReLU(nn.Module):
    """
    https://github.com/CoinCheung/DeepLab-v3-plus-cityscapes/blob/master/models/deeplabv3plus.py
    """

    def __init__(
        self, in_chan, out_chan, ks=3, stride=1, padding=1, dilation=1, *args, **kwargs
    ):
        super().__init__()
        self.conv = nn.Conv2d(
            in_chan,
            out_chan,
            kernel_size=ks,
            stride=stride,
            padding=padding,
            dilation=dilation,
            bias=True,
        )
        self.bn = nn.BatchNorm2d(out_chan)
        self.init_weight()

    def forward(self, x):
        x = self.conv(x)
        x = self.bn(x)
        return x

    def init_weight(self):
        for ly in self.children():
            if isinstance(ly, nn.Conv2d):
                nn.init.kaiming_normal_(ly.weight, a=1)
                if ly.bias is not None:
                    nn.init.constant_(ly.bias, 0)


class ASPPv3Plus(nn.Module):
    """
    https://github.com/CoinCheung/DeepLab-v3-plus-cityscapes/blob/master/models/deeplabv3plus.py
    """

    def __init__(self, backbone, no_init):
        super().__init__()

        if backbone == "mobilenet":
            in_chan = 320
        else:
            in_chan = 2048

        self.with_gp = False
        self.conv1 = ConvBNReLU(in_chan, 256, ks=1, dilation=1, padding=0)
        self.conv2 = ConvBNReLU(in_chan, 256, ks=3, dilation=6, padding=6)
        self.conv3 = ConvBNReLU(in_chan, 256, ks=3, dilation=12, padding=12)
        self.conv4 = ConvBNReLU(in_chan, 256, ks=3, dilation=18, padding=18)
        if self.with_gp:
            self.avg = nn.AdaptiveAvgPool2d((1, 1))
            self.conv1x1 = ConvBNReLU(in_chan, 256, ks=1)
            self.conv_out = ConvBNReLU(256 * 5, 256, ks=1)
        else:
            self.conv_out = ConvBNReLU(256 * 4, 256, ks=1)

        if not no_init:
            self.init_weight()

    def forward(self, x):
        H, W = x.size()[2:]
        feat1 = self.conv1(x)
        feat2 = self.conv2(x)
        feat3 = self.conv3(x)
        feat4 = self.conv4(x)
        if self.with_gp:
            avg = self.avg(x)
            feat5 = self.conv1x1(avg)
            feat5 = F.interpolate(feat5, (H, W), mode="bilinear", align_corners=True)
            feat = torch.cat([feat1, feat2, feat3, feat4, feat5], 1)
        else:
            feat = torch.cat([feat1, feat2, feat3, feat4], 1)
        feat = self.conv_out(feat)
        return feat

    def init_weight(self):
        for ly in self.children():
            if isinstance(ly, nn.Conv2d):
                nn.init.kaiming_normal_(ly.weight, a=1)
                if ly.bias is not None:
                    nn.init.constant_(ly.bias, 0)


class Decoder(nn.Module):
    """
    https://github.com/CoinCheung/DeepLab-v3-plus-cityscapes/blob/master/models/deeplabv3plus.py
    """

    def __init__(self, n_classes):
        super(Decoder, self).__init__()
        self.conv_low = ConvBNReLU(256, 48, ks=1, padding=0)
        self.conv_cat = nn.Sequential(
            ConvBNReLU(304, 256, ks=3, padding=1),
            ConvBNReLU(256, 256, ks=3, padding=1),
        )
        self.conv_out = nn.Conv2d(256, n_classes, kernel_size=1, bias=False)

    def forward(self, feat_low, feat_aspp):
        H, W = feat_low.size()[2:]
        feat_low = self.conv_low(feat_low)
        feat_aspp_up = F.interpolate(
            feat_aspp, (H, W), mode="bilinear", align_corners=True
        )
        feat_cat = torch.cat([feat_low, feat_aspp_up], dim=1)
        feat_out = self.conv_cat(feat_cat)
        logits = self.conv_out(feat_out)
        return logits


"""
https://github.com/jfzhang95/pytorch-deeplab-xception/blob/master/modeling/deeplab.py
"""


class DeepLabV3Decoder(nn.Module):
    def __init__(
        self,
        opts,
        no_init=False,
        freeze_bn=False,
    ):
        super().__init__()

        num_classes = opts.gen.s.output_dim
        self.backbone = opts.gen.deeplabv3.backbone
        self.use_dada = ("d" in opts.tasks) and opts.gen.s.use_dada

        if self.backbone == "resnet":
            self.aspp = ASPPv3Plus(self.backbone, no_init)
            self.decoder = Decoder(num_classes)

            self.freeze_bn = freeze_bn
        else:
            self.head = _DeepLabHead(num_classes, c4_channels=320)

        self._target_size = find_target_size(opts, "s")
        print(
            "      - {}:  setting target size to {}".format(
                self.__class__.__name__, self._target_size
            )
        )

        if not no_init:
            for m in self.modules():
                if isinstance(m, nn.Conv2d):
                    nn.init.kaiming_normal_(m.weight, mode="fan_out")
                    if m.bias is not None:
                        nn.init.zeros_(m.bias)
                elif isinstance(m, nn.BatchNorm2d):
                    nn.init.ones_(m.weight)
                    nn.init.zeros_(m.bias)
                elif isinstance(m, nn.Linear):
                    nn.init.normal_(m.weight, 0, 0.01)
                    nn.init.zeros_(m.bias)

            self.load_pretrained(opts)

    def load_pretrained(self, opts):
        assert opts.gen.deeplabv3.backbone in {"resnet", "mobilenet"}
        assert Path(opts.gen.deeplabv3.pretrained_model.resnet).exists()
        if opts.gen.deeplabv3.backbone == "resnet":
            std = torch.load(opts.gen.deeplabv3.pretrained_model.resnet)
            self.aspp.load_state_dict(
                {
                    k.replace("aspp.", ""): v
                    for k, v in std.items()
                    if k.startswith("aspp.")
                }
            )
            self.decoder.load_state_dict(
                {
                    k.replace("decoder.", ""): v
                    for k, v in std.items()
                    if k.startswith("decoder.")
                    and not (len(v.shape) > 0 and v.shape[0] == 19)
                },
                strict=False,
            )
            print(
                "- Loaded pre-trained DeepLabv3+ (Resnet) Decoder & ASPP as Seg Decoder"
            )
        else:
            std = torch.load(opts.gen.deeplabv3.pretrained_model.mobilenet)
            self.load_state_dict(
                {
                    k: v
                    for k, v in std.items()
                    if k.startswith("head.")
                    and not (len(v.shape) > 0 and v.shape[0] == 19)
                },
                strict=False,
            )
            print(
                "    - Loaded pre-trained DeepLabv3+ (MobileNetV2) Head as Seg Decoder"
            )

    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):
        assert isinstance(z, (tuple, list))
        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 ValueError(error)

        z_high, z_low = z

        if z_depth is not None and self.use_dada:
            z_high = z_high * z_depth

        if self.backbone == "resnet":
            z_high = self.aspp(z_high)
            s = self.decoder(z_high, z_low)
        else:
            s = self.head(z_high)

        s = F.interpolate(
            s, size=self._target_size, mode="bilinear", align_corners=True
        )

        return s

    def freeze_bn(self):
        for m in self.modules():
            if isinstance(m, nn.BatchNorm2d):
                m.eval()