| |
| |
|
|
| 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)) |
|
|
