def conv2d_ABN(ni, nf, stride, activation="leaky_relu", kernel_size=3, activation_param=1e-2, groups=1):
return nn.Sequential(
nn.Conv2d(ni, nf, kernel_size=kernel_size, stride=stride, padding=kernel_size // 2, groups=groups,
bias=False),
InPlaceABN(num_features=nf, activation=activation, activation_param=activation_param)
)
def __init__(self, in_feature_shape):
"""
Args:
- in_feature_shape (List[int]) : size of feature at different levels
"""
super().__init__()
# Channel information are the one given in the EfficientPS paper
# Depending on the EfficientNet model chosen the number of channel will
# change
# x4 size [B, 40, H, W] (input 40 channels)
# Bottom up path layers
self.conv_b_up_x4 = Conv2d(in_feature_shape[1], 256, 1)
self.iabn_b_up_x4 = InPlaceABN(256)
# Top down path layers
self.conv_t_dn_x4 = Conv2d(in_feature_shape[1], 256, 1)
self.iabn_t_dn_x4 = InPlaceABN(256)
# x8 size [B, 64, H, W] (input 64 channels)
# Bottom up path layers
self.conv_b_up_x8 = Conv2d(in_feature_shape[2], 256, 1)
self.iabn_b_up_x8 = InPlaceABN(256)
# Top down path layers
self.conv_t_dn_x8 = Conv2d(in_feature_shape[2], 256, 1)
self.iabn_t_dn_x8 = InPlaceABN(256)
# x16 size [B, 176, H, W] (input 176 channels)
# In the paper they took the 5 block of efficient net ie 128 channels
# But taking last block seem more pertinent and was already implemented
# Skipping to id 3 since block 4 does not interest us
# Bottom up path layers
self.conv_b_up_x16 = Conv2d(in_feature_shape[3], 256, 1)
self.iabn_b_up_x16 = InPlaceABN(256)
# Top down path layers
self.conv_t_dn_x16 = Conv2d(in_feature_shape[3], 256, 1)
self.iabn_t_dn_x16 = InPlaceABN(256)
# x32 size [B, 2048, H, W] (input 2048 channels)
# Bottom up path layers
self.conv_b_up_x32 = Conv2d(in_feature_shape[4], 256, 1)
self.iabn_b_up_x32 = InPlaceABN(256)
# Top down path layers
self.conv_t_dn_x32 = Conv2d(in_feature_shape[4], 256, 1)
self.iabn_t_dn_x32 = InPlaceABN(256)
# Separable Conv and Inplace BN at the output of the FPN
# x4
self.depth_wise_conv_x4 = DepthwiseSeparableConv(in_channels=256,
out_channels=256,
kernel_size=3,
padding=1)
self.iabn_out_x4 = InPlaceABN(256)
# x8
self.depth_wise_conv_x8 = DepthwiseSeparableConv(in_channels=256,
out_channels=256,
kernel_size=3,
padding=1)
self.iabn_out_x8 = InPlaceABN(256)
# x16
self.depth_wise_conv_x16 = DepthwiseSeparableConv(in_channels=256,
out_channels=256,
kernel_size=3,
padding=1)
self.iabn_out_x16 = InPlaceABN(256)
# x32
self.depth_wise_conv_x32 = DepthwiseSeparableConv(in_channels=256,
out_channels=256,
kernel_size=3,
padding=1)
self.iabn_out_x32 = InPlaceABN(256)