#
# For licensing see accompanying LICENSE file.
# Copyright (C) 2023 Apple Inc. All Rights Reserved.
#
import argparse
from typing import Optional, Tuple
import torch
import torch.nn.functional as F
from torch import Tensor, nn
from cvnets.layers import (
AdaptiveAvgPool2d,
BaseLayer,
ConvLayer2d,
Dropout2d,
SeparableConv2d,
)
from cvnets.modules import BaseModule
from utils import logger
from utils.ddp_utils import is_master
[docs]class ASPP(BaseModule):
"""
ASPP module defined in DeepLab papers, `here <https://arxiv.org/abs/1606.00915>`_ and `here <https://arxiv.org/abs/1706.05587>`_
Args:
opts: command-line arguments
in_channels (int): :math:`C_{in}` from an expected input of size :math:`(N, C_{in}, H, W)`
out_channels (int): :math:`C_{out}` from an expected output of size :math:`(N, C_{out}, H, W)`
atrous_rates (Tuple[int]): atrous rates for different branches.
is_sep_conv (Optional[bool]): Use separable convolution instead of standaard conv. Default: False
dropout (Optional[float]): Apply dropout. Default is 0.0
Shape:
- Input: :math:`(N, C_{in}, H, W)`
- Output: :math:`(N, C_{out}, H, W)`
"""
[docs] def __init__(
self,
opts,
in_channels: int,
out_channels: int,
atrous_rates: Tuple[int],
is_sep_conv: Optional[bool] = False,
dropout: Optional[float] = 0.0,
*args,
**kwargs
) -> None:
in_proj = ConvLayer2d(
opts=opts,
in_channels=in_channels,
out_channels=out_channels,
kernel_size=1,
stride=1,
use_norm=True,
use_act=True,
)
out_proj = ConvLayer2d(
opts=opts,
in_channels=5 * out_channels,
out_channels=out_channels,
kernel_size=1,
stride=1,
use_norm=True,
use_act=True,
)
aspp_layer = ASPPSeparableConv2d if is_sep_conv else ASPPConv2d
assert len(atrous_rates) == 3
modules = [in_proj]
modules.extend(
[
aspp_layer(
opts=opts,
in_channels=in_channels,
out_channels=out_channels,
dilation=rate,
)
for rate in atrous_rates
]
)
modules.append(
ASPPPooling(opts=opts, in_channels=in_channels, out_channels=out_channels)
)
if not (0.0 <= dropout < 1.0):
if is_master(opts):
logger.warning(
"Dropout value in {} should be between 0 and 1. Got: {}. Setting it to 0.0".format(
self.__class__.__name__, dropout
)
)
dropout = 0.0
super().__init__()
self.convs = nn.ModuleList(modules)
self.project = out_proj
self.in_channels = in_channels
self.out_channels = out_channels
self.atrous_rates = atrous_rates
self.is_sep_conv_layer = is_sep_conv
self.n_atrous_branches = len(atrous_rates)
self.dropout_layer = Dropout2d(p=dropout)
[docs] def forward(self, x: Tensor, *args, **kwargs) -> Tensor:
out = []
for conv in self.convs:
out.append(conv(x))
out = torch.cat(out, dim=1)
out = self.project(out)
out = self.dropout_layer(out)
return out
def __repr__(self):
return "{}(in_channels={}, out_channels={}, atrous_rates={}, is_aspp_sep={}, dropout={})".format(
self.__class__.__name__,
self.in_channels,
self.out_channels,
self.atrous_rates,
self.is_sep_conv_layer,
self.dropout_layer.p,
)
[docs]class ASPPConv2d(ConvLayer2d):
"""
Convolution with a dilation for the ASPP module
Args:
opts: command-line arguments
in_channels (int): :math:`C_{in}` from an expected input of size :math:`(N, C_{in}, H, W)`
out_channels (int): :math:`C_{out}` from an expected output of size :math:`(N, C_{out}, H, W)`
dilation (int): Dilation rate
Shape:
- Input: :math:`(N, C_{in}, H, W)`
- Output: :math:`(N, C_{out}, H, W)`
"""
[docs] def __init__(
self, opts, in_channels: int, out_channels: int, dilation: int, *args, **kwargs
) -> None:
super().__init__(
opts=opts,
in_channels=in_channels,
out_channels=out_channels,
kernel_size=3,
stride=1,
use_norm=True,
use_act=True,
dilation=dilation,
)
[docs] def adjust_atrous_rate(self, rate: int) -> None:
"""This function allows to adjust the dilation rate"""
self.block.conv.dilation = rate
# padding is the same here
# see ConvLayer to see the method for computing padding
self.block.conv.padding = rate
[docs]class ASPPSeparableConv2d(SeparableConv2d):
"""
Separable Convolution with a dilation for the ASPP module
Args:
opts: command-line arguments
in_channels (int): :math:`C_{in}` from an expected input of size :math:`(N, C_{in}, H, W)`
out_channels (int): :math:`C_{out}` from an expected output of size :math:`(N, C_{out}, H, W)`
dilation (int): Dilation rate
Shape:
- Input: :math:`(N, C_{in}, H, W)`
- Output: :math:`(N, C_{out}, H, W)`
"""
[docs] def __init__(
self, opts, in_channels: int, out_channels: int, dilation: int, *args, **kwargs
) -> None:
super().__init__(
opts=opts,
in_channels=in_channels,
out_channels=out_channels,
kernel_size=3,
stride=1,
dilation=dilation,
use_norm=True,
use_act=True,
)
[docs] def adjust_atrous_rate(self, rate: int) -> None:
"""This function allows to adjust the dilation rate"""
self.dw_conv.block.conv.dilation = rate
# padding is the same here
# see ConvLayer to see the method for computing padding
self.dw_conv.block.conv.padding = rate
[docs]class ASPPPooling(BaseLayer):
"""
ASPP pooling layer
Args:
opts: command-line arguments
in_channels (int): :math:`C_{in}` from an expected input of size :math:`(N, C_{in}, H, W)`
out_channels (int): :math:`C_{out}` from an expected output of size :math:`(N, C_{out}, H, W)`
Shape:
- Input: :math:`(N, C_{in}, H, W)`
- Output: :math:`(N, C_{out}, H, W)`
"""
[docs] def __init__(
self, opts, in_channels: int, out_channels: int, *args, **kwargs
) -> None:
super().__init__()
self.aspp_pool = nn.Sequential()
self.aspp_pool.add_module(
name="global_pool", module=AdaptiveAvgPool2d(output_size=1)
)
self.aspp_pool.add_module(
name="conv_1x1",
module=ConvLayer2d(
opts=opts,
in_channels=in_channels,
out_channels=out_channels,
kernel_size=1,
stride=1,
use_norm=True,
use_act=True,
),
)
self.in_channels = in_channels
self.out_channels = out_channels
[docs] def forward(self, x: Tensor) -> Tensor:
x_size = x.shape[-2:]
x = self.aspp_pool(x)
x = F.interpolate(x, size=x_size, mode="bilinear", align_corners=False)
return x
def __repr__(self):
return "{}(in_channels={}, out_channels={})".format(
self.__class__.__name__, self.in_channels, self.out_channels
)