File size: 3,024 Bytes
938e515 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 |
# -*- coding: utf-8 -*-
# Copyright (c) Facebook, Inc. and its affiliates.
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,
)
# default initialization
weight_init.c2_msra_fill(self.depthwise)
weight_init.c2_msra_fill(self.pointwise)
def forward(self, x):
return self.pointwise(self.depthwise(x))
|