|
|
|
|
|
|
|
import fvcore.nn.weight_init as weight_init
|
|
from torch import nn
|
|
|
|
from .batch_norm import FrozenBatchNorm2d, get_norm
|
|
from .wrappers import Conv2d
|
|
|
|
|
|
"""
|
|
CNN building blocks.
|
|
"""
|
|
|
|
|
|
class CNNBlockBase(nn.Module):
|
|
"""
|
|
A CNN block is assumed to have input channels, output channels and a stride.
|
|
The input and output of `forward()` method must be NCHW tensors.
|
|
The method can perform arbitrary computation but must match the given
|
|
channels and stride specification.
|
|
|
|
Attribute:
|
|
in_channels (int):
|
|
out_channels (int):
|
|
stride (int):
|
|
"""
|
|
|
|
def __init__(self, in_channels, out_channels, stride):
|
|
"""
|
|
The `__init__` method of any subclass should also contain these arguments.
|
|
|
|
Args:
|
|
in_channels (int):
|
|
out_channels (int):
|
|
stride (int):
|
|
"""
|
|
super().__init__()
|
|
self.in_channels = in_channels
|
|
self.out_channels = out_channels
|
|
self.stride = stride
|
|
|
|
def freeze(self):
|
|
"""
|
|
Make this block not trainable.
|
|
This method sets all parameters to `requires_grad=False`,
|
|
and convert all BatchNorm layers to FrozenBatchNorm
|
|
|
|
Returns:
|
|
the block itself
|
|
"""
|
|
for p in self.parameters():
|
|
p.requires_grad = False
|
|
FrozenBatchNorm2d.convert_frozen_batchnorm(self)
|
|
return self
|
|
|
|
|
|
class DepthwiseSeparableConv2d(nn.Module):
|
|
"""
|
|
A kxk depthwise convolution + a 1x1 convolution.
|
|
|
|
In :paper:`xception`, norm & activation are applied on the second conv.
|
|
:paper:`mobilenet` uses norm & activation on both convs.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
in_channels,
|
|
out_channels,
|
|
kernel_size=3,
|
|
padding=1,
|
|
dilation=1,
|
|
*,
|
|
norm1=None,
|
|
activation1=None,
|
|
norm2=None,
|
|
activation2=None,
|
|
):
|
|
"""
|
|
Args:
|
|
norm1, norm2 (str or callable): normalization for the two conv layers.
|
|
activation1, activation2 (callable(Tensor) -> Tensor): activation
|
|
function for the two conv layers.
|
|
"""
|
|
super().__init__()
|
|
self.depthwise = Conv2d(
|
|
in_channels,
|
|
in_channels,
|
|
kernel_size=kernel_size,
|
|
padding=padding,
|
|
dilation=dilation,
|
|
groups=in_channels,
|
|
bias=not norm1,
|
|
norm=get_norm(norm1, in_channels),
|
|
activation=activation1,
|
|
)
|
|
self.pointwise = Conv2d(
|
|
in_channels,
|
|
out_channels,
|
|
kernel_size=1,
|
|
bias=not norm2,
|
|
norm=get_norm(norm2, out_channels),
|
|
activation=activation2,
|
|
)
|
|
|
|
|
|
weight_init.c2_msra_fill(self.depthwise)
|
|
weight_init.c2_msra_fill(self.pointwise)
|
|
|
|
def forward(self, x):
|
|
return self.pointwise(self.depthwise(x))
|
|
|
