def __init__(self,
num_classes,
in_channels=3,
level2_depth=5,
level3_depth=3):
super().__init__()
self.level1 = layers.ConvBNPReLU(in_channels,
16,
3,
padding='same',
stride=2)
self.br1 = BNPReLU(in_channels + 16)
self.proj1 = layers.ConvBNPReLU(in_channels + 16, num_classes, 1)
self.level2_0 = DownSampler(in_channels + 16, 64)
self.level2 = nn.Sequential(
*[DilatedResidualBlock(64, 64) for i in range(level2_depth)])
self.br2 = BNPReLU(in_channels + 128)
self.proj2 = layers.ConvBNPReLU(in_channels + 128, num_classes, 1)
self.level3_0 = DownSampler(in_channels + 128, 128)
self.level3 = nn.Sequential(
*[DilatedResidualBlock(128, 128) for i in range(level3_depth)])
self.br3 = BNPReLU(256)
self.proj3 = layers.ConvBNPReLU(256, num_classes, 1)
def __init__(self,
num_classes,
in_channels=3,
scale=1.0,
drop_prob=0.1,
pretrained=None):
super().__init__()
self.backbone = EESPNetBackbone(in_channels, drop_prob, scale)
self.in_channels = self.backbone.out_channels
self.proj_l4_c = layers.ConvBNPReLU(self.in_channels[3],
self.in_channels[2],
1,
stride=1,
bias_attr=False)
psp_size = 2 * self.in_channels[2]
self.eesp_psp = nn.Sequential(
EESP(psp_size,
psp_size // 2,
stride=1,
branches=4,
kernel_size_maximum=7),
PSPModule(psp_size // 2, psp_size // 2),
)
self.project_l3 = nn.Sequential(
nn.Dropout2D(p=drop_prob),
nn.Conv2D(psp_size // 2, num_classes, 1, 1, bias_attr=False),
)
self.act_l3 = BNPReLU(num_classes)
self.project_l2 = layers.ConvBNPReLU(self.in_channels[1] + num_classes,
num_classes,
1,
stride=1,
bias_attr=False)
self.project_l1 = nn.Sequential(
nn.Dropout2D(p=drop_prob),
nn.Conv2D(self.in_channels[0] + num_classes,
num_classes,
1,
1,
bias_attr=False),
)
self.pretrained = pretrained
self.init_weight()
def __init__(self,
num_classes,
in_channels=3,
level2_depth=2,
level3_depth=3,
pretrained=None):
super().__init__()
self.encoder = ESPNetEncoder(num_classes, in_channels, level2_depth,
level3_depth)
self.level3_up = nn.Conv2DTranspose(num_classes,
num_classes,
2,
stride=2,
padding=0,
output_padding=0,
bias_attr=False)
self.br3 = layers.SyncBatchNorm(num_classes)
self.level2_proj = nn.Conv2D(in_channels + 128,
num_classes,
1,
bias_attr=False)
self.combine_l2_l3 = nn.Sequential(
BNPReLU(2 * num_classes),
DilatedResidualBlock(2 * num_classes, num_classes, residual=False),
)
self.level2_up = nn.Sequential(
nn.Conv2DTranspose(num_classes,
num_classes,
2,
stride=2,
padding=0,
output_padding=0,
bias_attr=False),
BNPReLU(num_classes),
)
self.out_proj = layers.ConvBNPReLU(16 + in_channels + num_classes,
num_classes,
3,
padding='same',
stride=1)
self.out_up = nn.Conv2DTranspose(num_classes,
num_classes,
2,
stride=2,
padding=0,
output_padding=0,
bias_attr=False)
self.pretrained = pretrained
def __init__(self, in_channels, out_channels, sizes=4):
super().__init__()
self.stages = nn.LayerList([
nn.Conv2D(in_channels,
in_channels,
kernel_size=3,
stride=1,
groups=in_channels,
padding='same',
bias_attr=False) for _ in range(sizes)
])
self.project = layers.ConvBNPReLU(in_channels * (sizes + 1),
out_channels,
1,
stride=1,
bias_attr=False)
def __init__(self,
in_channels,
out_channels,
branches=4,
kernel_size_maximum=9,
shortcut=True):
super().__init__()
if out_channels < in_channels:
raise RuntimeError(
"The out_channes for DownSampler should be bigger than in_channels, but got in_channles={}, out_channels={}"
.format(in_channels, out_channels))
self.eesp = EESP(in_channels,
out_channels - in_channels,
stride=2,
branches=branches,
kernel_size_maximum=kernel_size_maximum,
down_method='avg')
self.avg = nn.AvgPool2D(kernel_size=3, padding=1, stride=2)
if shortcut:
self.shortcut_layer = nn.Sequential(
layers.ConvBNPReLU(3, 3, 3, stride=1, bias_attr=False),
layers.ConvBN(3, out_channels, 1, stride=1, bias_attr=False),
)
self._act = nn.PReLU()
def __init__(self, in_channels=3, drop_prob=0.1, scale=1.0):
super().__init__()
reps = [0, 3, 7, 3]
num_level = 4 # 1/2, 1/4, 1/8, 1/16
kernel_size_limitations = [13, 11, 9, 7] # kernel size limitation
branch_list = [4] * len(
kernel_size_limitations) # branches at different levels
base_channels = 32 # first conv output channels
channels_config = [base_channels] * num_level
for i in range(num_level):
if i == 0:
channels = int(base_channels * scale)
channels = math.ceil(channels / branch_list[0]) * branch_list[0]
channels_config[
i] = base_channels if channels > base_channels else channels
else:
channels_config[i] = channels * pow(2, i)
self.level1 = layers.ConvBNPReLU(in_channels,
channels_config[0],
3,
stride=2,
bias_attr=False)
self.level2 = DownSampler(
channels_config[0],
channels_config[1],
branches=branch_list[0],
kernel_size_maximum=kernel_size_limitations[0],
shortcut=True)
self.level3_0 = DownSampler(
channels_config[1],
channels_config[2],
branches=branch_list[1],
kernel_size_maximum=kernel_size_limitations[1],
shortcut=True)
self.level3 = nn.LayerList()
for i in range(reps[1]):
self.level3.append(
EESP(channels_config[2],
channels_config[2],
stride=1,
branches=branch_list[2],
kernel_size_maximum=kernel_size_limitations[2]))
self.level4_0 = DownSampler(
channels_config[2],
channels_config[3],
branches=branch_list[2],
kernel_size_maximum=kernel_size_limitations[2],
shortcut=True)
self.level4 = nn.LayerList()
for i in range(reps[2]):
self.level4.append(
EESP(channels_config[3],
channels_config[3],
stride=1,
branches=branch_list[3],
kernel_size_maximum=kernel_size_limitations[3]))
self.out_channels = channels_config
self.init_params()
def __init__(self,
in_channels,
out_channels,
stride=1,
branches=4,
kernel_size_maximum=7,
down_method='esp'):
super(EESP, self).__init__()
if out_channels % branches != 0:
raise RuntimeError(
"The out_channes for EESP should be factorized by branches, but out_channels={} cann't be factorized by branches={}"
.format(out_channels, branches))
assert down_method in [
'avg', 'esp'
], "The down_method for EESP only support 'avg' or 'esp', but got down_method={}".format(
down_method)
self.in_channels = in_channels
self.stride = stride
in_branch_channels = int(out_channels / branches)
self.group_conv_in = layers.ConvBNPReLU(in_channels,
in_branch_channels,
1,
stride=1,
groups=branches,
bias_attr=False)
map_ksize_dilation = {
3: 1,
5: 2,
7: 3,
9: 4,
11: 5,
13: 6,
15: 7,
17: 8
}
self.kernel_sizes = []
for i in range(branches):
kernel_size = 3 + 2 * i
kernel_size = kernel_size if kernel_size <= kernel_size_maximum else 3
self.kernel_sizes.append(kernel_size)
self.kernel_sizes.sort()
self.spp_modules = nn.LayerList()
for i in range(branches):
dilation = map_ksize_dilation[self.kernel_sizes[i]]
self.spp_modules.append(
nn.Conv2D(in_branch_channels,
in_branch_channels,
kernel_size=3,
padding='same',
stride=stride,
dilation=dilation,
groups=in_branch_channels,
bias_attr=False))
self.group_conv_out = layers.ConvBN(out_channels,
out_channels,
kernel_size=1,
stride=1,
groups=branches,
bias_attr=False)
self.bn_act = BNPReLU(out_channels)
self._act = nn.PReLU()
self.down_method = True if down_method == 'avg' else False