def __init__(self, in_channels, out_channels=None, norm=True, activation=True, act_func='relu', onnx_export=True):
super(SeparableConvBlock, self).__init__()
if out_channels is None:
out_channels = in_channels
# Q: whether separate conv
# share bias between depthwise_conv and pointwise_conv
# or just pointwise_conv apply bias.
# A: Confirmed, just pointwise_conv applies bias, depthwise_conv has no bias.
# self.depthwise_conv = Conv2dStaticSamePadding(in_channels, in_channels,
# kernel_size=3, stride=1, groups=in_channels, bias=False)
# self.pointwise_conv = Conv2dStaticSamePadding(in_channels, out_channels, kernel_size=1, stride=1)
self.depthwise_conv = nn.Conv2d(in_channels, in_channels, kernel_size=3, stride=1, padding=1, bias=False, groups=in_channels)
self.pointwise_conv = nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=1, padding=0)
self.norm = norm
if self.norm:
# Warning: pytorch momentum is different from tensorflow's, momentum_pytorch = 1 - momentum_tensorflow
self.bn = nn.BatchNorm2d(num_features=out_channels, momentum=0.01, eps=1e-3)
self.activation = activation
if self.activation:
if act_func == 'swish':
self.act_func = MemoryEfficientSwish() if not onnx_export else Swish()
elif act_func == 'relu':
self.act_func = nn.ReLU()
else:
assert 0, f'Unknown act_func: {act_func}'
def __init__(self, in_channels, num_anchors, num_layers, onnx_export=False):
super(Regressor, self).__init__()
self.num_layers = num_layers
self.conv_list = nn.ModuleList(
[SeparableConvBlock(in_channels, in_channels, norm=False, activation=False) for i in range(num_layers)])
self.bn_list = nn.ModuleList(
[nn.ModuleList([nn.BatchNorm2d(in_channels, momentum=0.01, eps=1e-3) for i in range(num_layers)]) for j in
range(5)])
self.header = SeparableConvBlock(in_channels, num_anchors * 4, norm=False, activation=False)
self.swish = MemoryEfficientSwish() if not onnx_export else Swish()
def __init__(self, backbone_features_num, in_pyramid_levels=[3,4,5], out_pyramid_levels=[3,4,5,6,7], features_num=64,
first_time=False, epsilon=1e-4, onnx_export=True, attention=False, act_func='relu', **kwargs):
"""
Args:
num_channels:
conv_channels:
first_time: whether the input comes directly from the efficientnet,
if True, downchannel it first, and downsample P5 to generate P6 then P7
epsilon: epsilon of fast weighted attention sum of BiFPN, not the BN's epsilon
onnx_export: if True, use Swish instead of MemoryEfficientSwish
"""
assert in_pyramid_levels == [3,4,5]
assert out_pyramid_levels == [3,4,5,6,7]
assert act_func in ['relu', 'swish'], f'Unknown act_func: {act_func}'
super(BiFPN, self).__init__()
self.convert_onnx = False
self.pyramid_sizes = None
self.backbone_features_num = backbone_features_num
self.in_pyramid_levels = in_pyramid_levels
self.out_pyramid_levels = out_pyramid_levels
self.features_num = features_num
self.epsilon = epsilon
self.attention = attention
logger = kwargs.get('logger', None)
if logger:
logger.info('Build Neck: BiFPN')
logger.info(f'Features Num: {features_num}')
logger.info(f'First: {first_time}')
logger.info(f'Attention: {attention}')
# Conv layers
self.conv3_up = SeparableConvBlock(features_num, activation=True, act_func=act_func, onnx_export=onnx_export)
self.conv4_up = SeparableConvBlock(features_num, activation=True, act_func=act_func, onnx_export=onnx_export)
self.conv5_up = SeparableConvBlock(features_num, activation=True, act_func=act_func, onnx_export=onnx_export)
self.conv6_up = SeparableConvBlock(features_num, activation=True, act_func=act_func, onnx_export=onnx_export)
self.conv7_down = SeparableConvBlock(features_num, activation=True, act_func=act_func, onnx_export=onnx_export)
self.conv6_down = SeparableConvBlock(features_num, activation=True, act_func=act_func, onnx_export=onnx_export)
self.conv5_down = SeparableConvBlock(features_num, activation=True, act_func=act_func, onnx_export=onnx_export)
self.conv4_down = SeparableConvBlock(features_num, activation=True, act_func=act_func, onnx_export=onnx_export)
# print(self.conv6_up)
# print(self.conv4_down)
# Feature scaling layers
self.p6_upsample = nn.Upsample(scale_factor=2, mode='nearest')
self.p5_upsample = nn.Upsample(scale_factor=2, mode='nearest')
self.p4_upsample = nn.Upsample(scale_factor=2, mode='nearest')
self.p3_upsample = nn.Upsample(scale_factor=2, mode='nearest')
# self.p4_downsample = MaxPool2dStaticSamePadding(3, 2)
self.p4_downsample = nn.MaxPool2d(3, stride=2, padding=1)
self.p5_downsample = nn.MaxPool2d(3, stride=2, padding=1)
self.p6_downsample = nn.MaxPool2d(3, stride=2, padding=1)
self.p7_downsample = nn.MaxPool2d(3, stride=2, padding=1)
# self.swish = MemoryEfficientSwish() if not onnx_export else Swish()
if act_func == 'swish':
self.act_func = MemoryEfficientSwish() if not onnx_export else Swish()
elif act_func == 'relu':
self.act_func = nn.ReLU()
self.first_time = first_time
if self.first_time:
self.p5_down_channel = nn.Sequential(
# Conv2dStaticSamePadding(conv_channels[2], features_num, 1),
# nn.BatchNorm2d(features_num, momentum=0.01, eps=1e-3),
nn.Conv2d(backbone_features_num[2], features_num, 1),
nn.BatchNorm2d(features_num, momentum=0.01, eps=1e-3),
)
self.p4_down_channel = nn.Sequential(
nn.Conv2d(backbone_features_num[1], features_num, 1),
nn.BatchNorm2d(features_num, momentum=0.01, eps=1e-3),
)
self.p3_down_channel = nn.Sequential(
nn.Conv2d(backbone_features_num[0], features_num, 1),
nn.BatchNorm2d(features_num, momentum=0.01, eps=1e-3),
)
self.p5_to_p6 = nn.Sequential(
nn.Conv2d(backbone_features_num[2], features_num, 1),
nn.BatchNorm2d(features_num, momentum=0.01, eps=1e-3),
nn.MaxPool2d(3, stride=2, padding=1)
)
self.p6_to_p7 = nn.Sequential(
nn.MaxPool2d(3, stride=2, padding=1)
)
self.p4_down_channel_2 = nn.Sequential(
nn.Conv2d(backbone_features_num[1], features_num, 1),
nn.BatchNorm2d(features_num, momentum=0.01, eps=1e-3),
)
self.p5_down_channel_2 = nn.Sequential(
nn.Conv2d(backbone_features_num[2], features_num, 1),
nn.BatchNorm2d(features_num, momentum=0.01, eps=1e-3),
)
# Weight
if self.attention:
self.p6_w1 = nn.Parameter(torch.ones(2, dtype=torch.float32), requires_grad=True)
self.p6_w1_relu = nn.ReLU()
self.p5_w1 = nn.Parameter(torch.ones(2, dtype=torch.float32), requires_grad=True)
self.p5_w1_relu = nn.ReLU()
self.p4_w1 = nn.Parameter(torch.ones(2, dtype=torch.float32), requires_grad=True)
self.p4_w1_relu = nn.ReLU()
self.p3_w1 = nn.Parameter(torch.ones(2, dtype=torch.float32), requires_grad=True)
self.p3_w1_relu = nn.ReLU()
self.p4_w2 = nn.Parameter(torch.ones(3, dtype=torch.float32), requires_grad=True)
self.p4_w2_relu = nn.ReLU()
self.p5_w2 = nn.Parameter(torch.ones(3, dtype=torch.float32), requires_grad=True)
self.p5_w2_relu = nn.ReLU()
self.p6_w2 = nn.Parameter(torch.ones(3, dtype=torch.float32), requires_grad=True)
self.p6_w2_relu = nn.ReLU()
self.p7_w2 = nn.Parameter(torch.ones(2, dtype=torch.float32), requires_grad=True)
self.p7_w2_relu = nn.ReLU()
self._initialize_weights(logger=logger)
def __init__(self, backbone_features_num, in_pyramid_levels=[3,4,5], neck_pyramid_levels=[3,4,5,6,7],
out_pyramid_levels=[3,4,5,6,7],
features_num=64, attention=False, act_func='relu',
index=-1, first_time=False, last_time=False, epsilon=1e-4, onnx_export=True, **kwargs):
"""
Args:
num_channels:
conv_channels:
first_time: whether the input comes directly from the efficientnet,
if True, downchannel it first, and downsample P5 to generate P6 then P7
epsilon: epsilon of fast weighted attention sum of BiFPN, not the BN's epsilon
onnx_export: if True, use Swish instead of MemoryEfficientSwish
"""
assert act_func in ['relu', 'swish'], f'Unknown act_func: {act_func}'
higher_pyramid_levels_num = len(neck_pyramid_levels) - len(in_pyramid_levels)
assert higher_pyramid_levels_num <= 2, 'Higher pyramid levels num is too large.'
assert higher_pyramid_levels_num >= 1, 'Higher pyramid levels num is too small.'
super(BiFPNModule, self).__init__()
self.convert_onnx = False
self.pyramid_sizes = None
self.backbone_features_num = backbone_features_num
self.in_pyramid_levels = in_pyramid_levels
self.neck_pyramid_levels = neck_pyramid_levels
self.out_pyramid_levels = out_pyramid_levels
self.features_num = features_num
self.attention = attention
self.index = index
self.first_time = first_time
self.last_time = last_time
self.epsilon = epsilon
self.higher_pyramid_levels_num = higher_pyramid_levels_num
in_pyramid_num = len(in_pyramid_levels)
neck_pyramid_num = len(neck_pyramid_levels)
out_pyramid_num = len(out_pyramid_levels)
self.in_pyramid_num = in_pyramid_num
self.neck_pyramid_num = neck_pyramid_num
self.out_pyramid_num = out_pyramid_num
self.out_neck_pyramid_idx = list()
for idx, p in enumerate(neck_pyramid_levels):
if p in out_pyramid_levels:
self.out_neck_pyramid_idx.append(idx)
logger = kwargs.get('logger', None)
if logger:
logger.info(f'---- Build BiFPN Module [{index}] ----')
logger.info(f'Features Num: {features_num}')
logger.info(f'Attention: {attention}')
logger.info(f'First : {first_time}')
logger.info(f'Last : {last_time}')
logger.info(f'In Pyramid Levels : {in_pyramid_levels} ( Num: {in_pyramid_num})')
logger.info(f'Neck Pyramid Levels : {neck_pyramid_levels} ( Num: {neck_pyramid_num})')
logger.info(f'Out Pyramid Levels : {out_pyramid_levels} ( Num: {out_pyramid_num})')
logger.info(f'Out Neck Pyramid Idx : {self.out_neck_pyramid_idx}')
# Conv layers
self.conv_up = nn.ModuleList()
for i in range(neck_pyramid_num-1):
self.conv_up.append(
SeparableConvBlock(features_num, activation=True, act_func=act_func, onnx_export=onnx_export)
)
self.conv_down = nn.ModuleList()
_conv_down_num = neck_pyramid_num-1 if not last_time else out_pyramid_num-1
for i in range(_conv_down_num):
self.conv_down.append(
SeparableConvBlock(features_num, activation=True, act_func=act_func, onnx_export=onnx_export)
)
# Feature scaling layers
self.up_sample = nn.ModuleList()
for i in range(neck_pyramid_num-1):
self.up_sample.append(nn.Upsample(scale_factor=2, mode='nearest'))
self.down_sample = nn.MaxPool2d(3, stride=2, padding=1)
# self.swish = MemoryEfficientSwish() if not onnx_export else Swish()
if act_func == 'swish':
self.act_func = MemoryEfficientSwish() if not onnx_export else Swish()
elif act_func == 'relu':
self.act_func = nn.ReLU()
if self.first_time:
self.down_channel = nn.ModuleList()
for i in range(in_pyramid_num):
self.down_channel.append(nn.Sequential(
nn.Conv2d(backbone_features_num[i], features_num, 1),
nn.BatchNorm2d(features_num, momentum=0.01, eps=1e-3),
))
if higher_pyramid_levels_num >= 1:
self.higher_pyramid_conv = nn.ModuleList()
self.higher_pyramid_conv.append(
nn.Sequential(
nn.Conv2d(backbone_features_num[2], features_num, 1),
nn.BatchNorm2d(features_num, momentum=0.01, eps=1e-3),
nn.MaxPool2d(3, stride=2, padding=1)))
if higher_pyramid_levels_num >= 2:
self.higher_pyramid_conv.append(
nn.MaxPool2d(3, stride=2, padding=1))
#self.down_channel_2 = nn.ModuleList()
#for i in range(in_pyramid_num - 1):
# self.down_channel_2.append(
# nn.Sequential(
# nn.Conv2d(backbone_features_num[i+1], features_num, 1),
# nn.BatchNorm2d(features_num, momentum=0.01, eps=1e-3),
# )
# )
self.down_channel_two = nn.ModuleList()
for i in range(in_pyramid_num - 1):
self.down_channel_two.append(
nn.Sequential(
nn.Conv2d(backbone_features_num[i+1], features_num, 1),
nn.BatchNorm2d(features_num, momentum=0.01, eps=1e-3),
)
)
# Weight
if self.attention:
# assert 0, 'not support now'
self.relu = nn.ReLU()
self.fuse_w_up = nn.ParameterList()
# p(0) ~ p(N-2)
for i in range(neck_pyramid_num-1):
self.fuse_w_up.append(nn.Parameter(torch.ones(2, dtype=torch.float32), requires_grad=True))
self.fuse_w_down = nn.ParameterList()
# p(1) ~ p(N-2)
for i in range(1, neck_pyramid_num-1):
self.fuse_w_down.append(nn.Parameter(torch.ones(3, dtype=torch.float32), requires_grad=True))
# p(N-1)
self.fuse_w_down.append(nn.Parameter(torch.ones(2, dtype=torch.float32), requires_grad=True))
self._initialize_weights(logger=logger)
def __init__(self, num_channels, conv_channels, first_time=False, epsilon=1e-4, onnx_export=False, attention=True):
"""
Args:
num_channels:
conv_channels:
first_time: whether the input comes directly from the efficientnet,
if True, downchannel it first, and downsample P5 to generate P6 then P7
epsilon: epsilon of fast weighted attention sum of BiFPN, not the BN's epsilon
onnx_export: if True, use Swish instead of MemoryEfficientSwish
"""
super(BiFPN, self).__init__()
self.epsilon = epsilon
# Conv layers
self.conv6_up = SeparableConvBlock(num_channels, onnx_export=onnx_export)
self.conv5_up = SeparableConvBlock(num_channels, onnx_export=onnx_export)
self.conv4_up = SeparableConvBlock(num_channels, onnx_export=onnx_export)
self.conv3_up = SeparableConvBlock(num_channels, onnx_export=onnx_export)
self.conv4_down = SeparableConvBlock(num_channels, onnx_export=onnx_export)
self.conv5_down = SeparableConvBlock(num_channels, onnx_export=onnx_export)
self.conv6_down = SeparableConvBlock(num_channels, onnx_export=onnx_export)
self.conv7_down = SeparableConvBlock(num_channels, onnx_export=onnx_export)
# Feature scaling layers
self.p6_upsample = nn.Upsample(scale_factor=2, mode='nearest')
self.p5_upsample = nn.Upsample(scale_factor=2, mode='nearest')
self.p4_upsample = nn.Upsample(scale_factor=2, mode='nearest')
self.p3_upsample = nn.Upsample(scale_factor=2, mode='nearest')
self.p4_downsample = MaxPool2dStaticSamePadding(3, 2)
self.p5_downsample = MaxPool2dStaticSamePadding(3, 2)
self.p6_downsample = MaxPool2dStaticSamePadding(3, 2)
self.p7_downsample = MaxPool2dStaticSamePadding(3, 2)
self.swish = MemoryEfficientSwish() if not onnx_export else Swish()
self.first_time = first_time
if self.first_time:
self.p5_down_channel = nn.Sequential(
Conv2dStaticSamePadding(conv_channels[2], num_channels, 1),
nn.BatchNorm2d(num_channels, momentum=0.01, eps=1e-3),
)
self.p4_down_channel = nn.Sequential(
Conv2dStaticSamePadding(conv_channels[1], num_channels, 1),
nn.BatchNorm2d(num_channels, momentum=0.01, eps=1e-3),
)
self.p3_down_channel = nn.Sequential(
Conv2dStaticSamePadding(conv_channels[0], num_channels, 1),
nn.BatchNorm2d(num_channels, momentum=0.01, eps=1e-3),
)
self.p5_to_p6 = nn.Sequential(
Conv2dStaticSamePadding(conv_channels[2], num_channels, 1),
nn.BatchNorm2d(num_channels, momentum=0.01, eps=1e-3),
MaxPool2dStaticSamePadding(3, 2)
)
self.p6_to_p7 = nn.Sequential(
MaxPool2dStaticSamePadding(3, 2)
)
self.p4_down_channel_2 = nn.Sequential(
Conv2dStaticSamePadding(conv_channels[1], num_channels, 1),
nn.BatchNorm2d(num_channels, momentum=0.01, eps=1e-3),
)
self.p5_down_channel_2 = nn.Sequential(
Conv2dStaticSamePadding(conv_channels[2], num_channels, 1),
nn.BatchNorm2d(num_channels, momentum=0.01, eps=1e-3),
)
# Weight
self.p6_w1 = nn.Parameter(torch.ones(2, dtype=torch.float32), requires_grad=True)
self.p6_w1_relu = nn.ReLU()
self.p5_w1 = nn.Parameter(torch.ones(2, dtype=torch.float32), requires_grad=True)
self.p5_w1_relu = nn.ReLU()
self.p4_w1 = nn.Parameter(torch.ones(2, dtype=torch.float32), requires_grad=True)
self.p4_w1_relu = nn.ReLU()
self.p3_w1 = nn.Parameter(torch.ones(2, dtype=torch.float32), requires_grad=True)
self.p3_w1_relu = nn.ReLU()
self.p4_w2 = nn.Parameter(torch.ones(3, dtype=torch.float32), requires_grad=True)
self.p4_w2_relu = nn.ReLU()
self.p5_w2 = nn.Parameter(torch.ones(3, dtype=torch.float32), requires_grad=True)
self.p5_w2_relu = nn.ReLU()
self.p6_w2 = nn.Parameter(torch.ones(3, dtype=torch.float32), requires_grad=True)
self.p6_w2_relu = nn.ReLU()
self.p7_w2 = nn.Parameter(torch.ones(2, dtype=torch.float32), requires_grad=True)
self.p7_w2_relu = nn.ReLU()
self.attention = attention