def __init__(self, in_channels: int, out_channels: int, name=None):
super(TransitionLayer, self).__init__()
num_in = len(in_channels)
num_out = len(out_channels)
self.conv_bn_func_list = []
for i in range(num_out):
residual = None
if i < num_in:
if in_channels[i] != out_channels[i]:
residual = self.add_sublayer(
"transition_{}_layer_{}".format(name, i + 1),
L.ConvBNReLU(in_channels=in_channels[i],
out_channels=out_channels[i],
kernel_size=3,
padding='same',
bias_attr=False))
else:
residual = self.add_sublayer(
"transition_{}_layer_{}".format(name, i + 1),
L.ConvBNReLU(in_channels=in_channels[-1],
out_channels=out_channels[i],
kernel_size=3,
stride=2,
padding='same',
bias_attr=False))
self.conv_bn_func_list.append(residual)
def __init__(self,
num_channels: int,
num_filters: int,
has_se: bool,
stride: int = 1,
downsample: bool = False,
name: str = None):
super(BottleneckBlock, self).__init__()
self.has_se = has_se
self.downsample = downsample
self.conv1 = L.ConvBNReLU(in_channels=num_channels,
out_channels=num_filters,
kernel_size=1,
padding='same',
bias_attr=False)
self.conv2 = L.ConvBNReLU(in_channels=num_filters,
out_channels=num_filters,
kernel_size=3,
stride=stride,
padding='same',
bias_attr=False)
self.conv3 = L.ConvBN(in_channels=num_filters,
out_channels=num_filters * 4,
kernel_size=1,
padding='same',
bias_attr=False)
if self.downsample:
self.conv_down = L.ConvBN(in_channels=num_channels,
out_channels=num_filters * 4,
kernel_size=1,
padding='same',
bias_attr=False)
if self.has_se:
self.se = SELayer(num_channels=num_filters * 4,
num_filters=num_filters * 4,
reduction_ratio=16,
name=name + '_fc')
def __init__(self,
in_channels: int,
key_channels: int,
out_channels: int,
dropout_rate: float = 0.1):
super().__init__()
self.attention_block = ObjectAttentionBlock(in_channels, key_channels)
self.conv1x1 = nn.Sequential(
L.ConvBNReLU(2 * in_channels, out_channels, 1),
nn.Dropout2D(dropout_rate))
def __init__(self,
in_channels: int,
out_channels: int,
multi_scale_output: bool = True,
name: str = None,
align_corners: bool = False):
super(FuseLayers, self).__init__()
self._actual_ch = len(in_channels) if multi_scale_output else 1
self._in_channels = in_channels
self.align_corners = align_corners
self.residual_func_list = []
for i in range(self._actual_ch):
for j in range(len(in_channels)):
if j > i:
residual_func = self.add_sublayer(
"residual_{}_layer_{}_{}".format(name, i + 1, j + 1),
L.ConvBN(in_channels=in_channels[j],
out_channels=out_channels[i],
kernel_size=1,
padding='same',
bias_attr=False))
self.residual_func_list.append(residual_func)
elif j < i:
pre_num_filters = in_channels[j]
for k in range(i - j):
if k == i - j - 1:
residual_func = self.add_sublayer(
"residual_{}_layer_{}_{}_{}".format(
name, i + 1, j + 1, k + 1),
L.ConvBN(in_channels=pre_num_filters,
out_channels=out_channels[i],
kernel_size=3,
stride=2,
padding='same',
bias_attr=False))
pre_num_filters = out_channels[i]
else:
residual_func = self.add_sublayer(
"residual_{}_layer_{}_{}_{}".format(
name, i + 1, j + 1, k + 1),
L.ConvBNReLU(in_channels=pre_num_filters,
out_channels=out_channels[j],
kernel_size=3,
stride=2,
padding='same',
bias_attr=False))
pre_num_filters = out_channels[j]
self.residual_func_list.append(residual_func)
def __init__(self,
num_classes: int,
in_channels: int,
ocr_mid_channels: int = 512,
ocr_key_channels: int = 256):
super().__init__()
self.num_classes = num_classes
self.spatial_gather = SpatialGatherBlock()
self.spatial_ocr = SpatialOCRModule(ocr_mid_channels, ocr_key_channels,
ocr_mid_channels)
self.indices = [-2, -1] if len(in_channels) > 1 else [-1, -1]
self.conv3x3_ocr = L.ConvBNReLU(in_channels[self.indices[1]],
ocr_mid_channels,
3,
padding=1)
self.cls_head = nn.Conv2D(ocr_mid_channels, self.num_classes, 1)
self.aux_head = nn.Sequential(
L.ConvBNReLU(in_channels[self.indices[0]],
in_channels[self.indices[0]], 1),
nn.Conv2D(in_channels[self.indices[0]], self.num_classes, 1))
def __init__(self, in_channels: int, key_channels: int):
super().__init__()
self.in_channels = in_channels
self.key_channels = key_channels
self.f_pixel = nn.Sequential(
L.ConvBNReLU(in_channels, key_channels, 1),
L.ConvBNReLU(key_channels, key_channels, 1))
self.f_object = nn.Sequential(
L.ConvBNReLU(in_channels, key_channels, 1),
L.ConvBNReLU(key_channels, key_channels, 1))
self.f_down = L.ConvBNReLU(in_channels, key_channels, 1)
self.f_up = L.ConvBNReLU(key_channels, in_channels, 1)
def __init__(self,
pretrained: str = None,
stage1_num_modules: int = 1,
stage1_num_blocks: tuple = (4, ),
stage1_num_channels: tuple = (64, ),
stage2_num_modules: int = 1,
stage2_num_blocks: tuple = (4, 4),
stage2_num_channels: tuple = (18, 36),
stage3_num_modules: int = 4,
stage3_num_blocks: tuple = (4, 4, 4),
stage3_num_channels: tuple = (18, 36, 72),
stage4_num_modules: int = 3,
stage4_num_blocks: tuple = (4, 4, 4, 4),
stage4_num_channels: tuple = (18, 36, 72, 144),
has_se: bool = False,
align_corners: bool = False):
super(HRNet_W18, self).__init__()
self.pretrained = pretrained
self.stage1_num_modules = stage1_num_modules
self.stage1_num_blocks = stage1_num_blocks
self.stage1_num_channels = stage1_num_channels
self.stage2_num_modules = stage2_num_modules
self.stage2_num_blocks = stage2_num_blocks
self.stage2_num_channels = stage2_num_channels
self.stage3_num_modules = stage3_num_modules
self.stage3_num_blocks = stage3_num_blocks
self.stage3_num_channels = stage3_num_channels
self.stage4_num_modules = stage4_num_modules
self.stage4_num_blocks = stage4_num_blocks
self.stage4_num_channels = stage4_num_channels
self.has_se = has_se
self.align_corners = align_corners
self.feat_channels = [sum(stage4_num_channels)]
self.conv_layer1_1 = L.ConvBNReLU(in_channels=3,
out_channels=64,
kernel_size=3,
stride=2,
padding='same',
bias_attr=False)
self.conv_layer1_2 = L.ConvBNReLU(in_channels=64,
out_channels=64,
kernel_size=3,
stride=2,
padding='same',
bias_attr=False)
self.la1 = Layer1(num_channels=64,
num_blocks=self.stage1_num_blocks[0],
num_filters=self.stage1_num_channels[0],
has_se=has_se,
name="layer2")
self.tr1 = TransitionLayer(
in_channels=[self.stage1_num_channels[0] * 4],
out_channels=self.stage2_num_channels,
name="tr1")
self.st2 = Stage(num_channels=self.stage2_num_channels,
num_modules=self.stage2_num_modules,
num_blocks=self.stage2_num_blocks,
num_filters=self.stage2_num_channels,
has_se=self.has_se,
name="st2",
align_corners=align_corners)
self.tr2 = TransitionLayer(in_channels=self.stage2_num_channels,
out_channels=self.stage3_num_channels,
name="tr2")
self.st3 = Stage(num_channels=self.stage3_num_channels,
num_modules=self.stage3_num_modules,
num_blocks=self.stage3_num_blocks,
num_filters=self.stage3_num_channels,
has_se=self.has_se,
name="st3",
align_corners=align_corners)
self.tr3 = TransitionLayer(in_channels=self.stage3_num_channels,
out_channels=self.stage4_num_channels,
name="tr3")
self.st4 = Stage(num_channels=self.stage4_num_channels,
num_modules=self.stage4_num_modules,
num_blocks=self.stage4_num_blocks,
num_filters=self.stage4_num_channels,
has_se=self.has_se,
name="st4",
align_corners=align_corners)