def __init__(self,
in_channels,
num_anchors,
num_classes,
num_layers,
pyramid_levels=5,
onnx_export=False):
super(Classifier, self).__init__()
self.num_anchors = num_anchors
self.num_classes = num_classes
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(pyramid_levels)
])
self.header = SeparableConvBlock(in_channels,
num_anchors * num_classes,
norm=False,
activation=False)
self.swish = MemoryEfficientSwish() if not onnx_export else Swish()
def __init__(self, block_args, global_params):
super().__init__()
self._block_args = block_args
self._bn_mom = 1 - global_params.batch_norm_momentum
self._bn_eps = global_params.batch_norm_epsilon
self.has_se = (self._block_args.se_ratio
is not None) and (0 < self._block_args.se_ratio <= 1)
self.id_skip = block_args.id_skip # skip connection and drop connect
# Get static or dynamic convolution depending on image size
Conv2d = get_same_padding_conv2d(image_size=global_params.image_size)
# Expansion phase
inp = self._block_args.input_filters # number of input channels
oup = self._block_args.input_filters * self._block_args.expand_ratio # number of output channels
if self._block_args.expand_ratio != 1:
self._expand_conv = Conv2d(in_channels=inp,
out_channels=oup,
kernel_size=1,
bias=False)
self._bn0 = nn.BatchNorm2d(num_features=oup,
momentum=self._bn_mom,
eps=self._bn_eps)
# Depthwise convolution phase
k = self._block_args.kernel_size
s = self._block_args.stride
self._depthwise_conv = Conv2d(
in_channels=oup,
out_channels=oup,
groups=oup, # groups makes it depthwise
kernel_size=k,
stride=s,
bias=False)
self._bn1 = nn.BatchNorm2d(num_features=oup,
momentum=self._bn_mom,
eps=self._bn_eps)
# Squeeze and Excitation layer, if desired
if self.has_se:
num_squeezed_channels = max(
1,
int(self._block_args.input_filters *
self._block_args.se_ratio))
self._se_reduce = Conv2d(in_channels=oup,
out_channels=num_squeezed_channels,
kernel_size=1)
self._se_expand = Conv2d(in_channels=num_squeezed_channels,
out_channels=oup,
kernel_size=1)
# Output phase
final_oup = self._block_args.output_filters
self._project_conv = Conv2d(in_channels=oup,
out_channels=final_oup,
kernel_size=1,
bias=False)
self._bn2 = nn.BatchNorm2d(num_features=final_oup,
momentum=self._bn_mom,
eps=self._bn_eps)
self._swish = MemoryEfficientSwish()
def __init__(self, blocks_args=None, global_params=None):
super().__init__()
assert isinstance(blocks_args, list), 'blocks_args should be a list'
assert len(blocks_args) > 0, 'block args must be greater than 0'
self._global_params = global_params
self._blocks_args = blocks_args
# Get static or dynamic convolution depending on image size
Conv2d = get_same_padding_conv2d(image_size=global_params.image_size)
# Batch norm parameters
bn_mom = 1 - self._global_params.batch_norm_momentum
bn_eps = self._global_params.batch_norm_epsilon
# Stem
in_channels = 3 # rgb
out_channels = round_filters(32, self._global_params) # number of output channels
self._conv_stem = Conv2d(in_channels, out_channels, kernel_size=3, stride=2, bias=False)
self._bn0 = nn.BatchNorm2d(num_features=out_channels, momentum=bn_mom, eps=bn_eps)
# Build blocks
self._blocks = nn.ModuleList([])
for i in range(len(self._blocks_args)):
# Update block input and output filters based on depth multiplier.
self._blocks_args[i] = self._blocks_args[i]._replace(
input_filters=round_filters(self._blocks_args[i].input_filters, self._global_params),
output_filters=round_filters(self._blocks_args[i].output_filters, self._global_params),
num_repeat=round_repeats(self._blocks_args[i].num_repeat, self._global_params)
)
# The first block needs to take care of stride and filter size increase.
self._blocks.append(MBConvBlock(self._blocks_args[i], self._global_params))
if self._blocks_args[i].num_repeat > 1:
self._blocks_args[i] = self._blocks_args[i]._replace(input_filters=self._blocks_args[i].output_filters, stride=1)
for _ in range(self._blocks_args[i].num_repeat - 1):
self._blocks.append(MBConvBlock(self._blocks_args[i], self._global_params))
# Head'efficientdet-d0': 'efficientnet-b0',
in_channels = self._blocks_args[len(self._blocks_args)-1].output_filters # output of final block
out_channels = round_filters(1280, self._global_params)
self._conv_head = Conv2d(in_channels, out_channels, kernel_size=1, bias=False)
self._bn1 = nn.BatchNorm2d(num_features=out_channels, momentum=bn_mom, eps=bn_eps)
# Final linear layer
self._avg_pooling = nn.AdaptiveAvgPool2d(1)
self._dropout = nn.Dropout(self._global_params.dropout_rate)
self._fc = nn.Linear(out_channels, self._global_params.num_classes)
self._swish = MemoryEfficientSwish()
def set_swish(self, memory_efficient=True):
"""Sets swish function as memory efficient (for training) or standard (for export)"""
self._swish = MemoryEfficientSwish() if memory_efficient else Swish()
for block in self._blocks:
block.set_swish(memory_efficient)
def __init__(self,
num_channels,
conv_channels,
first_time=False,
epsilon=1e-4,
onnx_export=False,
attention=True,
use_p8=False):
"""
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
self.use_p8 = use_p8
# 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)
if use_p8:
self.conv7_up = SeparableConvBlock(num_channels,
onnx_export=onnx_export)
self.conv8_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)
if use_p8:
self.p7_upsample = nn.Upsample(scale_factor=2, mode='nearest')
self.p8_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))
if use_p8:
self.p7_to_p8 = 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
