def __init__(self,
in_channels,
num_anchors,
num_classes,
num_layers,
onnx_export=False):
super(ClassNet, 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)
])
# 每一个有效特征层对应的BatchNorm2d不同
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)
])
# num_anchors = 9
# num_anchors num_classes
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,
in_channels,
out_channels=None,
norm=True,
activation=False,
onnx_export=False):
super(SeparableConvBlock, self).__init__()
if out_channels is None:
out_channels = in_channels
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.norm = norm
if self.norm:
self.bn = nn.BatchNorm2d(num_features=out_channels,
momentum=0.01,
eps=1e-3)
self.activation = activation
if self.activation:
self.swish = MemoryEfficientSwish() if not onnx_export else Swish()
def __init__(self, block_args, global_params):
super().__init__()
self._block_args = block_args
# 获得一种卷积方法
Conv2d = get_same_padding_conv2d(image_size=global_params.image_size)
# 获得标准化的参数
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
#-------------------------------------------------#
# 利用Inverted residuals
# part1 利用1x1卷积进行通道数上升
#-------------------------------------------------#
inp = self._block_args.input_filters
oup = self._block_args.input_filters * self._block_args.expand_ratio
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)
#------------------------------------------------------#
# 如果步长为2x2的话，利用深度可分离卷积进行高宽压缩
# part2 利用3x3卷积对每一个channel进行卷积
#------------------------------------------------------#
k = self._block_args.kernel_size
s = self._block_args.stride
self._depthwise_conv = Conv2d(in_channels=oup, out_channels=oup, groups=oup, kernel_size=k, stride=s, bias=False)
self._bn1 = nn.BatchNorm2d(num_features=oup, momentum=self._bn_mom, eps=self._bn_eps)
#------------------------------------------------------#
# 完成深度可分离卷积后
# 对深度可分离卷积的结果施加注意力机制
#------------------------------------------------------#
if self.has_se:
num_squeezed_channels = max(1, int(self._block_args.input_filters * self._block_args.se_ratio))
#------------------------------------------------------#
# 通道先压缩后上升，最后利用sigmoid将值固定到0-1之间
#------------------------------------------------------#
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)
#------------------------------------------------------#
# part3 利用1x1卷积进行通道下降
#------------------------------------------------------#
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
# 获得一种卷积方法，固定了image_size
Conv2d = get_same_padding_conv2d(image_size=global_params.image_size)
# 获得标准化的参数
bn_mom = 1 - self._global_params.batch_norm_momentum
bn_eps = self._global_params.batch_norm_epsilon
# 网络主干部分开始
# 设定输入进来的是RGB三通道图像
in_channels = 3
# 利用round_filters可以使得通道数在扩张的时候可以被8整除
out_channels = round_filters(32, self._global_params)
# 卷积+标准化
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)
# 对每个block的参数进行修改
self._blocks = nn.ModuleList([])
for i in range(len(self._blocks_args)):
# 对每个block的参数进行修改，根据所选的efficient版本进行修改
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)
)
# 第一次大的Block里面的卷积需要考虑步长和输入进来的通道数！
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部分
in_channels = self._blocks_args[len(self._blocks_args) - 1].output_filters
out_channels = round_filters(1280, self._global_params) # 32倍降的输出通道
# 卷积+标准化
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)
# 最后的全连接层
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)
# 进行swish激活函数
self._swish = MemoryEfficientSwish()
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
Conv2d = get_same_padding_conv2d(image_size=global_params.image_size)
# 1x1卷积通道扩张
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)
# 深度可分离卷积
k = self._block_args.kernel_size
s = self._block_args.stride
self._depthwise_conv = Conv2d(
in_channels=oup, out_channels=oup, groups=oup,
kernel_size=k, stride=s, bias=False)
self._bn1 = nn.BatchNorm2d(
num_features=oup, momentum=self._bn_mom, eps=self._bn_eps)
# 注意力机制模块组，先进行通道数的收缩再进行通道数的扩张
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)
# 输出部分
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 set_swish(self, memory_efficient=True):
"""Sets swish function as memory efficient (for training) or standard (for export)"""
# swish函数
self._swish = MemoryEfficientSwish() if memory_efficient else Swish()
for block in self._blocks:
block.set_swish(memory_efficient)
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()
def __init__(self, num_channels, conv_channels, first_time=False, epsilon=1e-4, onnx_export=False, attention=True):
super(BiFPN, self).__init__()
self.epsilon = epsilon
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)
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:
# 获取到了efficientnet的最后三层，对其进行通道的下压缩
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),
)
# 对输入进来的p5进行宽高的下采样
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)
)
# BIFPN第一轮的时候，跳线那里并不是同一个in
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),
)
# 简易注意力机制的weights
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
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
# 获得一种卷积方法
Conv2d = get_same_padding_conv2d(image_size=global_params.image_size)
# 获得标准化的参数
bn_mom = 1 - self._global_params.batch_norm_momentum
bn_eps = self._global_params.batch_norm_epsilon
#-------------------------------------------------#
# 网络主干部分开始
# 设定输入进来的是RGB三通道图像
# 利用round_filters可以使得通道可以被8整除
#-------------------------------------------------#
in_channels = 3
out_channels = round_filters(32, self._global_params)
#-------------------------------------------------#
# 创建stem部分
#-------------------------------------------------#
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)
#-------------------------------------------------#
# 在这个地方对大结构块进行循环
#-------------------------------------------------#
self._blocks = nn.ModuleList([])
for i in range(len(self._blocks_args)):
#-------------------------------------------------------------#
# 对每个block的参数进行修改，根据所选的efficient版本进行修改
#-------------------------------------------------------------#
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)
)
#-------------------------------------------------------------#
# 每个大结构块里面的第一个EfficientBlock
# 都需要考虑步长和输入通道数
#-------------------------------------------------------------#
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)
#---------------------------------------------------------------#
# 在利用第一个EfficientBlock进行通道数的调整或者高和宽的压缩后
# 进行EfficientBlock的堆叠
#---------------------------------------------------------------#
for _ in range(self._blocks_args[i].num_repeat - 1):
self._blocks.append(MBConvBlock(self._blocks_args[i], self._global_params))
#----------------------------------------------------------------#
# 这是efficientnet的尾部部分，在进行effcientdet构建的时候没用到
# 只在利用efficientnet进行分类的时候用到。
#----------------------------------------------------------------#
in_channels = self._blocks_args[len(self._blocks_args)-1].output_filters
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)
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()
