def __init__(self,
nclass,
backbone='',
aux=False,
jpu=False,
pretrained_base=False,
**kwargs):
super(ESPNetV2, self).__init__()
self.pretrained = eespnet(pretrained=pretrained_base, **kwargs)
self.proj_L4_C = _ConvBNPReLU(256, 128, 1, **kwargs)
self.pspMod = nn.Sequential(
EESP(256, 128, stride=1, k=4, r_lim=7, **kwargs),
_PSPModule(128, 128, **kwargs))
self.project_l3 = nn.Sequential(nn.Dropout2d(0.1),
nn.Conv2d(128, nclass, 1, bias=False))
self.act_l3 = _BNPReLU(nclass, **kwargs)
self.project_l2 = _ConvBNPReLU(64 + nclass, nclass, 1, **kwargs)
self.project_l1 = nn.Sequential(
nn.Dropout2d(0.1), nn.Conv2d(32 + nclass, nclass, 1, bias=False))
self.aux = aux
self.__setattr__('exclusive', [
'proj_L4_C', 'pspMod', 'project_l3', 'act_l3', 'project_l2',
'project_l1'
])
def __init__(self,
nclass,
backbone='',
aux=False,
jpu=False,
pretrained_base=True,
M=3,
N=21,
**kwargs):
super(CGNet, self).__init__()
# stage 1
self.stage1_0 = _ConvBNPReLU(3, 32, 3, 2, 1, **kwargs)
self.stage1_1 = _ConvBNPReLU(32, 32, 3, 1, 1, **kwargs)
self.stage1_2 = _ConvBNPReLU(32, 32, 3, 1, 1, **kwargs)
self.sample1 = _InputInjection(1)
self.sample2 = _InputInjection(2)
self.bn_prelu1 = _BNPReLU(32 + 3, **kwargs)
# stage 2
self.stage2_0 = ContextGuidedBlock(32 + 3,
64,
dilation=2,
reduction=8,
down=True,
residual=False,
**kwargs)
self.stage2 = nn.ModuleList()
for i in range(0, M - 1):
self.stage2.append(
ContextGuidedBlock(64, 64, dilation=2, reduction=8, **kwargs))
self.bn_prelu2 = _BNPReLU(128 + 3, **kwargs)
# stage 3
self.stage3_0 = ContextGuidedBlock(128 + 3,
128,
dilation=4,
reduction=16,
down=True,
residual=False,
**kwargs)
self.stage3 = nn.ModuleList()
for i in range(0, N - 1):
self.stage3.append(
ContextGuidedBlock(128,
128,
dilation=4,
reduction=16,
**kwargs))
self.bn_prelu3 = _BNPReLU(256, **kwargs)
self.head = nn.Sequential(nn.Dropout2d(0.1, False),
nn.Conv2d(256, nclass, 1))
self.__setattr__('exclusive', [
'stage1_0', 'stage1_1', 'stage1_2', 'sample1', 'sample2',
'bn_prelu1', 'stage2_0', 'stage2', 'bn_prelu2', 'stage3_0',
'stage3', 'bn_prelu3', 'head'
])
def __init__(
self,
in_channels,
out_channels,
k=4,
r_lim=9,
reinf=True,
inp_reinf=3,
norm_layer=None,
):
super(DownSampler, self).__init__()
channels_diff = out_channels - in_channels
self.eesp = EESP(
in_channels,
channels_diff,
stride=2,
k=k,
r_lim=r_lim,
down_method="avg",
norm_layer=norm_layer,
)
self.avg = nn.AvgPool2d(kernel_size=3, padding=1, stride=2)
if reinf:
self.inp_reinf = nn.Sequential(
_ConvBNPReLU(inp_reinf, inp_reinf, 3, 1, 1),
_ConvBN(inp_reinf, out_channels, 1, 1),
)
self.act = nn.PReLU(out_channels)
def __init__(self,
in_channels,
out_channels,
stride=1,
k=4,
r_lim=7,
down_method='esp',
norm_layer=nn.BatchNorm2d):
super(EESP, self).__init__()
self.stride = stride
n = int(out_channels / k)
n1 = out_channels - (k - 1) * n
assert down_method in ['avg', 'esp'
], 'One of these is suppported (avg or esp)'
assert n == n1, "n(={}) and n1(={}) should be equal for Depth-wise Convolution ".format(
n, n1)
self.proj_1x1 = _ConvBNPReLU(in_channels,
n,
1,
stride=1,
groups=k,
norm_layer=norm_layer)
map_receptive_ksize = {
3: 1,
5: 2,
7: 3,
9: 4,
11: 5,
13: 6,
15: 7,
17: 8
}
self.k_sizes = list()
for i in range(k):
ksize = int(3 + 2 * i)
ksize = ksize if ksize <= r_lim else 3
self.k_sizes.append(ksize)
self.k_sizes.sort()
self.spp_dw = nn.ModuleList()
for i in range(k):
dilation = map_receptive_ksize[self.k_sizes[i]]
self.spp_dw.append(
nn.Conv2d(n,
n,
3,
stride,
dilation,
dilation=dilation,
groups=n,
bias=False))
self.conv_1x1_exp = _ConvBN(out_channels,
out_channels,
1,
1,
groups=k,
norm_layer=norm_layer)
self.br_after_cat = _BNPReLU(out_channels, norm_layer)
self.module_act = nn.PReLU(out_channels)
self.downAvg = True if down_method == 'avg' else False
def __init__(self,
in_channels,
out_channels,
dilation=2,
reduction=16,
down=False,
residual=True,
norm_layer=nn.BatchNorm2d,
**kwargs):
super(ContextGuidedBlock, self).__init__()
inter_channels = out_channels // 2 if not down else out_channels
if down:
self.conv = _ConvBNPReLU(in_channels,
inter_channels,
3,
2,
1,
norm_layer=norm_layer,
**kwargs)
self.reduce = nn.Conv2d(inter_channels * 2,
out_channels,
1,
bias=False)
else:
self.conv = _ConvBNPReLU(in_channels,
inter_channels,
1,
1,
0,
norm_layer=norm_layer,
**kwargs)
self.f_loc = _ChannelWiseConv(inter_channels, inter_channels, **kwargs)
self.f_sur = _ChannelWiseConv(inter_channels, inter_channels, dilation,
**kwargs)
self.bn = norm_layer(inter_channels * 2)
self.prelu = nn.PReLU(inter_channels * 2)
self.f_glo = _FGlo(out_channels, reduction, **kwargs)
self.down = down
self.residual = residual
def __init__(self,
in_channels,
out_channels=1024,
sizes=(1, 2, 4, 8),
**kwargs):
super(_PSPModule, self).__init__()
self.stages = nn.ModuleList([
nn.Conv2d(in_channels,
in_channels,
3,
1,
1,
groups=in_channels,
bias=False) for _ in sizes
])
self.project = _ConvBNPReLU(in_channels * (len(sizes) + 1),
out_channels, 1, 1, **kwargs)
def __init__(self,
num_classes=1000,
scale=1,
reinf=True,
norm_layer=nn.BatchNorm2d):
super(EESPNet, self).__init__()
inp_reinf = 3 if reinf else None
reps = [0, 3, 7, 3]
r_lim = [13, 11, 9, 7, 5]
K = [4] * len(r_lim)
# set out_channels
base, levels, base_s = 32, 5, 0
out_channels = [base] * levels
for i in range(levels):
if i == 0:
base_s = int(base * scale)
base_s = math.ceil(base_s / K[0]) * K[0]
out_channels[i] = base if base_s > base else base_s
else:
out_channels[i] = base_s * pow(2, i)
if scale <= 1.5:
out_channels.append(1024)
elif scale in [1.5, 2]:
out_channels.append(1280)
else:
raise ValueError("Unknown scale value.")
self.level1 = _ConvBNPReLU(3,
out_channels[0],
3,
2,
1,
norm_layer=norm_layer)
self.level2_0 = DownSampler(
out_channels[0],
out_channels[1],
k=K[0],
r_lim=r_lim[0],
reinf=reinf,
inp_reinf=inp_reinf,
norm_layer=norm_layer,
)
self.level3_0 = DownSampler(
out_channels[1],
out_channels[2],
k=K[1],
r_lim=r_lim[1],
reinf=reinf,
inp_reinf=inp_reinf,
norm_layer=norm_layer,
)
self.level3 = nn.ModuleList()
for i in range(reps[1]):
self.level3.append(
EESP(
out_channels[2],
out_channels[2],
k=K[2],
r_lim=r_lim[2],
norm_layer=norm_layer,
))
self.level4_0 = DownSampler(
out_channels[2],
out_channels[3],
k=K[2],
r_lim=r_lim[2],
reinf=reinf,
inp_reinf=inp_reinf,
norm_layer=norm_layer,
)
self.level4 = nn.ModuleList()
for i in range(reps[2]):
self.level4.append(
EESP(
out_channels[3],
out_channels[3],
k=K[3],
r_lim=r_lim[3],
norm_layer=norm_layer,
))
self.level5_0 = DownSampler(
out_channels[3],
out_channels[4],
k=K[3],
r_lim=r_lim[3],
reinf=reinf,
inp_reinf=inp_reinf,
norm_layer=norm_layer,
)
self.level5 = nn.ModuleList()
for i in range(reps[2]):
self.level5.append(
EESP(
out_channels[4],
out_channels[4],
k=K[4],
r_lim=r_lim[4],
norm_layer=norm_layer,
))
self.level5.append(
_ConvBNPReLU(
out_channels[4],
out_channels[4],
3,
1,
1,
groups=out_channels[4],
norm_layer=norm_layer,
))
self.level5.append(
_ConvBNPReLU(
out_channels[4],
out_channels[5],
1,
1,
0,
groups=K[4],
norm_layer=norm_layer,
))
self.fc = nn.Linear(out_channels[5], num_classes)
for m in self.modules():
if isinstance(m, nn.Conv2d):
nn.init.kaiming_normal_(m.weight,
mode="fan_out",
nonlinearity="relu")
if m.bias is not None:
nn.init.constant_(m.bias, 0)
elif isinstance(m, nn.BatchNorm2d):
nn.init.constant_(m.weight, 1)
nn.init.constant_(m.bias, 0)
elif isinstance(m, nn.Linear):
nn.init.normal_(m.weight, std=0.001)
if m.bias is not None:
nn.init.constant_(m.bias, 0)