def __init__(self):
super(OutdoorSceneSeg, self).__init__()
# conv1
blocks = []
conv1_1 = B.conv_block(3, 64, 3, 2, 1, 1, False, 'zero', 'batch') # /2
conv1_2 = B.conv_block(64, 64, 3, 1, 1, 1, False, 'zero', 'batch')
conv1_3 = B.conv_block(64, 128, 3, 1, 1, 1, False, 'zero', 'batch')
max_pool = nn.MaxPool2d(3, stride=2, padding=0, ceil_mode=True) # /2
blocks = [conv1_1, conv1_2, conv1_3, max_pool]
# conv2, 3 blocks
blocks.append(Res131(128, 64, 256))
for i in range(2):
blocks.append(Res131(256, 64, 256))
# conv3, 4 blocks
blocks.append(Res131(256, 128, 512, 1, 2)) # /2
for i in range(3):
blocks.append(Res131(512, 128, 512))
# conv4, 23 blocks
blocks.append(Res131(512, 256, 1024, 2))
for i in range(22):
blocks.append(Res131(1024, 256, 1024, 2))
# conv5
blocks.append(Res131(1024, 512, 2048, 4))
blocks.append(Res131(2048, 512, 2048, 4))
blocks.append(Res131(2048, 512, 2048, 4))
blocks.append(B.conv_block(2048, 512, 3, 1, 1, 1, False, 'zero', 'batch'))
blocks.append(nn.Dropout(0.1))
# # conv6
blocks.append(nn.Conv2d(512, 8, 1, 1))
self.feature = B.sequential(*blocks)
# deconv
self.deconv = nn.ConvTranspose2d(8, 8, 16, 8, 4, 0, 8, False, 1)
# softmax
self.softmax = nn.Softmax(1)
def __init__(self, in_nc, out_nc, nf, nb, gc=32, upscale=4, norm_type=None, act_type='leakyrelu', \
mode='CNA', res_scale=1, upsample_mode='upconv'):
super(RRDB_Net, self).__init__()
n_upscale = int(math.log(upscale, 2))
if upscale == 3:
n_upscale = 1
fea_conv = B.conv_block(in_nc, nf, kernel_size=3, norm_type=None, act_type=None)
rb_blocks = [B.RRDB(nf, kernel_size=3, gc=32, stride=1, bias=True, pad_type='zero', \
norm_type=norm_type, act_type=act_type, mode='CNA') for _ in range(nb)]
LR_conv = B.conv_block(nf, nf, kernel_size=3, norm_type=norm_type, act_type=None, mode=mode)
if upsample_mode == 'upconv':
upsample_block = B.upconv_blcok
elif upsample_mode == 'pixelshuffle':
upsample_block = B.pixelshuffle_block
else:
raise NotImplementedError('upsample mode [%s] is not found' % upsample_mode)
if upscale == 3:
upsampler = upsample_block(nf, nf, 3, act_type=act_type)
else:
upsampler = [upsample_block(nf, nf, act_type=act_type) for _ in range(n_upscale)]
HR_conv0 = B.conv_block(nf, nf, kernel_size=3, norm_type=None, act_type=act_type)
HR_conv1 = B.conv_block(nf, out_nc, kernel_size=3, norm_type=None, act_type=None)
self.model = B.sequential(fea_conv, B.ShortcutBlock(B.sequential(*rb_blocks, LR_conv)),\
*upsampler, HR_conv0, HR_conv1)
def __init__(self,
in_nc=5,
out_nc=4,
nf=64,
nb=None,
upscale=4,
norm_type=None,
act_type='relu',
mode='CNA',
upsample_mode='upconv'):
super().__init__()
self.upscale = upscale
conv1 = B.conv_block(in_nc,
nf,
kernel_size=9,
norm_type=None,
act_type=act_type)
conv2 = B.conv_block(nf,
nf // 2,
kernel_size=5,
norm_type=None,
act_type=act_type)
conv3 = B.conv_block(nf // 2,
out_nc,
kernel_size=5,
norm_type=None,
act_type=None)
self.model = B.sequential(conv1, conv2, conv3)
def __init__(self, in_nc=3, out_nc=3, nf=64, nb=16, upscale=4, norm_type= None , act_type='relu', \
mode='CNA', res_scale=1, upsample_mode='pixelshuffle'):
super(SRResNet, self).__init__()
n_upscale = int(math.log(upscale, 2))
if upscale == 3:
n_upscale = 1
fea_conv = B.conv_block(in_nc, nf, kernel_size=3, norm_type=None, act_type=None)
resnet_blocks = [B.ResNetBlock(nf, nf, nf, norm_type=norm_type, act_type=act_type,\
mode=mode, res_scale=res_scale) for _ in range(nb)]
LR_conv = B.conv_block(nf, nf, kernel_size=3, norm_type=norm_type, act_type=None, mode=mode)
if upsample_mode == 'upconv':
upsample_block = B.upconv_blcok
elif upsample_mode == 'pixelshuffle':
upsample_block = B.pixelshuffle_block
else:
raise NotImplementedError('upsample mode [{:s}] is not found'.format(upsample_mode))
if upscale == 3:
upsampler = upsample_block(nf, nf, 3, act_type=act_type)
else:
upsampler = [upsample_block(nf, nf, act_type=act_type) for _ in range(n_upscale)]
HR_conv0 = B.conv_block(nf, nf, kernel_size=3, norm_type=None, act_type=act_type)
HR_conv1 = B.conv_block(nf, out_nc, kernel_size=3, norm_type=None, act_type=None)
self.model = B.sequential(fea_conv, B.ShortcutBlock(B.sequential(*resnet_blocks, LR_conv)),\
*upsampler, HR_conv0, HR_conv1)
def make_upsample_block(upscale=4, in_ch=64, out_nc=3, kernel_size=3):
n_upscale = 1 if upscale == 3 else int(math.log(upscale, 2))
LR_conv = B.conv_block(nf,
nf,
kernel_size=3,
norm_type=None,
act_type=None,
mode='CNA')
HR_conv0 = B.conv_block(nf,
nf,
kernel_size=3,
norm_type=None,
act_type='leakyrelu')
HR_conv1 = B.conv_block(nf,
out_nc,
kernel_size=kernel_size,
norm_type=None,
act_type=None)
if upscale == 1:
return nn.Sequential(LR_conv, HR_conv0, HR_conv1)
elif upscale == 3:
upsampler = B.upconv_blcok(nf, nf, 3, act_type=act_type)
else:
upsampler = [
B.upconv_blcok(nf, nf, act_type=act_type)
for _ in range(n_upscale)
]
return nn.Sequential(LR_conv, *upsampler, HR_conv0, HR_conv1)
def __init__(self, in_nc, mid_nc, out_nc, dilation=1, stride=1):
super(Res131, self).__init__()
conv0 = B.conv_block(in_nc, mid_nc, 1, 1, 1, 1, False, 'zero', 'batch')
conv1 = B.conv_block(mid_nc, mid_nc, 3, stride, dilation, 1, False, 'zero', 'batch')
conv2 = B.conv_block(mid_nc, out_nc, 1, 1, 1, 1, False, 'zero', 'batch', None) # No ReLU
self.res = B.sequential(conv0, conv1, conv2)
if in_nc == out_nc:
self.has_proj = False
else:
self.has_proj = True
self.proj = B.conv_block(in_nc, out_nc, 1, stride, 1, 1, False, 'zero', 'batch', None)
def __init__(self, in_nc=1, out_nc=1, nf=16, nb=16, gc=16, norm_type=None, act_type='leakyrelu', mode='CNA'):
super(DenoseNet, self).__init__()
fea_conv = B.conv_block(in_nc, nf, kernel_size=3, norm_type=None, act_type=None)
rb_blocks = [B.RRDB(nf, kernel_size=3, gc=gc, stride=1, bias=True, pad_type='zero',
norm_type=norm_type, act_type=act_type, mode='CNA', dilation=(1, 2, 4), groups=1) for _ in
range(nb)]
LR_conv = B.conv_block(nf, nf, kernel_size=3, norm_type=norm_type, act_type=None, mode=mode)
HR_conv0 = B.conv_block(nf, nf, kernel_size=3, norm_type=None, act_type=act_type)
HR_conv1 = B.conv_block(nf, out_nc, kernel_size=3, norm_type=None, act_type=None)
self.model = B.sequential(fea_conv, B.ShortcutBlock(B.sequential(*rb_blocks, LR_conv))
, HR_conv0, HR_conv1, nn.Sigmoid())
def __init__(self,
in_nc,
out_nc,
nf,
nb,
gc=32,
upscale=4,
norm_type=None,
act_type='leakyrelu',
res_scale=1):
super(RRDB_Net, self).__init__()
n_upscale = int(math.log(upscale, 2))
fea_conv = B.conv_block(in_nc,
nf,
kernel_size=3,
norm_type=None,
act_type=None)
rb_blocks = [
B.RRDB(nf,
kernel_size=3,
gc=32,
stride=1,
bias=True,
norm_type=norm_type,
act_type=act_type) for _ in range(nb)
]
LR_conv = B.conv_block(nf,
nf,
kernel_size=3,
norm_type=norm_type,
act_type=None)
upsample_block = B.upconv_blcok
upsampler = [
upsample_block(nf, nf, act_type=act_type) for _ in range(n_upscale)
]
HR_conv0 = B.conv_block(nf,
nf,
kernel_size=3,
norm_type=None,
act_type=act_type)
HR_conv1 = B.conv_block(nf,
out_nc,
kernel_size=3,
norm_type=None,
act_type=None)
self.model = B.sequential(
fea_conv, B.ShortcutBlock(B.sequential(*rb_blocks, LR_conv)),
*upsampler, HR_conv0, HR_conv1)
def __init__(self,
lm_nc,
pan_nc=1,
downscale=4,
norm_type=None,
act_type='leakyrelu',
mode='CNA'):
super().__init__()
assert downscale in (2, 4), 'downscale should be 2 or 4'
self.downscale = downscale
# Rec LM
# conv_LM0 = B.conv_block(
# 1,
# 1,
# kernel_size=4,
# stride=2,
# norm_type=None,
# act_type=act_type,
# mode=mode) if downscale == 4 else None
# conv_LM1 = B.conv_block(
# 1,
# 1,
# kernel_size=4,
# stride=2,
# norm_type=None,
# act_type=None,
# mode=mode)
conv_LM0 = B.conv_block(lm_nc,
lm_nc,
kernel_size=3,
stride=1,
groups=lm_nc,
bias=False,
norm_type=None,
act_type=None,
mode=mode)
down_LM1 = nn.Upsample(scale_factor=1 / downscale, mode='bilinear')
self.LM = B.sequential(conv_LM0, down_LM1)
# Rec Pan
conv_PAN0 = B.conv_block(lm_nc,
pan_nc,
kernel_size=1,
stride=1,
bias=False,
norm_type=None,
act_type=None,
mode=mode)
self.PAN = B.sequential(conv_PAN0)
def __init__(self, in_nc=3, nf=50, num_modules=4, out_nc=3, upscale=4):
"""
Constructor, see class documentation for more details.
Parameters
----------
in_nc : int
The number of input channels to the network, by default 3
nf : int
The number of input feature channels per block of the network
num_modules : int
The number of blocks in the network, by default 4
out_nc : int
The number of channels which should be output by the network, by
default 3
upscale : int
The amount of upscaling to be done, by default 4
"""
super(RFDN, self).__init__()
self.fea_conv = B.conv_layer(in_nc, nf, kernel_size=3)
self.B1 = B.RFDB(in_channels=nf)
self.B2 = B.RFDB(in_channels=nf)
self.B3 = B.RFDB(in_channels=nf)
self.B4 = B.RFDB(in_channels=nf)
self.c = B.conv_block(nf * num_modules, nf, kernel_size=1, act_type='lrelu')
self.LR_conv = B.conv_layer(nf, nf, kernel_size=3)
upsample_block = B.pixelshuffle_block
self.upsampler = upsample_block(nf, out_nc, upscale_factor=upscale)
self.scale_idx = 0
def __init__(self,
in_nc,
base_nf,
norm_type='batch',
act_type='leakyrelu'):
super(Discriminator_VGG_96, self).__init__()
# features
# hxw, c
# 96, 64
conv0 = B.conv_block(in_nc,
base_nf,
kernel_size=3,
norm_type=None,
act_type=act_type)
conv1 = B.conv_block(base_nf, base_nf, kernel_size=4, stride=2, norm_type=norm_type, \
act_type=act_type)
# 48, 64
conv2 = B.conv_block(base_nf, base_nf*2, kernel_size=3, stride=1, norm_type=norm_type, \
act_type=act_type)
conv3 = B.conv_block(base_nf*2, base_nf*2, kernel_size=4, stride=2, norm_type=norm_type, \
act_type=act_type)
# 24, 128
conv4 = B.conv_block(base_nf*2, base_nf*4, kernel_size=3, stride=1, norm_type=norm_type, \
act_type=act_type)
conv5 = B.conv_block(base_nf*4, base_nf*4, kernel_size=4, stride=2, norm_type=norm_type, \
act_type=act_type)
# 12, 256
conv6 = B.conv_block(base_nf*4, base_nf*8, kernel_size=3, stride=1, norm_type=norm_type, \
act_type=act_type)
conv7 = B.conv_block(base_nf*8, base_nf*8, kernel_size=4, stride=2, norm_type=norm_type, \
act_type=act_type)
# 6, 512
conv8 = B.conv_block(base_nf*8, base_nf*8, kernel_size=3, stride=1, norm_type=norm_type, \
act_type=act_type)
conv9 = B.conv_block(base_nf*8, base_nf*8, kernel_size=4, stride=2, norm_type=norm_type, \
act_type=act_type)
# 3, 512
self.features = B.sequential(conv0, conv1, conv2, conv3, conv4, conv5, conv6, conv7, conv8,\
conv9)
# classifier
self.classifier = nn.Sequential(nn.Linear(512 * 3 * 3, 100),
nn.LeakyReLU(0.2, True),
nn.Linear(100, 1))
def hourglass(data, nFilters, nModules, n, workspace, name, binarize, dcn):
s = 2
_dcn = False
up1 = data
for i in xrange(nModules):
up1 = conv_block(up1, nFilters, (1,1), True, "%s_up1_%d"%(name,i), binarize, _dcn, 1)
low1 = mx.sym.Pooling(data=data, kernel=(s, s), stride=(s,s), pad=(0,0), pool_type='max')
for i in xrange(nModules):
low1 = conv_block(low1, nFilters, (1,1), True, "%s_low1_%d"%(name,i), binarize, _dcn, 1)
if n>1:
low2 = hourglass(low1, nFilters, nModules, n-1, workspace, "%s_%d"%(name, n-1), binarize, dcn)
else:
low2 = low1
for i in xrange(nModules):
low2 = conv_block(low2, nFilters, (1,1), True, "%s_low2_%d"%(name,i), binarize, _dcn, 1) #TODO
low3 = low2
for i in xrange(nModules):
low3 = conv_block(low3, nFilters, (1,1), True, "%s_low3_%d"%(name,i), binarize, _dcn, 1)
up2 = mx.symbol.UpSampling(low3, scale=s, sample_type='nearest', workspace=512, name='%s_upsampling_%s'%(name,n), num_args=1)
return mx.symbol.add_n(up1, up2)
def __init__(self, in_nc=3, nf=50, num_modules=4, out_nc=3, upscale=4):
super(RFDN, self).__init__()
self.fea_conv = B.conv_layer(in_nc, nf, kernel_size=3)
self.B1 = B.RFDB(in_channels=nf)
self.B2 = B.RFDB(in_channels=nf)
self.B3 = B.RFDB(in_channels=nf)
self.B4 = B.RFDB(in_channels=nf)
self.c = B.conv_block(nf * num_modules, nf, kernel_size=1, act_type='lrelu')
self.LR_conv = B.conv_layer(nf, nf, kernel_size=3)
upsample_block = B.pixelshuffle_block
self.upsampler = upsample_block(nf, out_nc, upscale_factor=4)
self.scale_idx = 0
def __init__(
self,
lm_nc=4,
# pan_nc=1,
base_nf=64,
# downscale=4,
norm_type='batch',
act_type='leakyrelu',
mode='CNA'):
super().__init__()
# features
# h and w, c
# input 256, 3
conv0 = B.conv_block(lm_nc,
base_nf,
kernel_size=3,
norm_type=None,
act_type=act_type,
mode=mode)
conv1 = B.conv_block(base_nf,
base_nf,
kernel_size=4,
stride=2,
norm_type=norm_type,
act_type=act_type,
mode=mode)
# 128, 64
conv2 = B.conv_block(base_nf,
base_nf * 2,
kernel_size=3,
stride=1,
norm_type=norm_type,
act_type=act_type,
mode=mode)
conv3 = B.conv_block(base_nf * 2,
base_nf * 2,
kernel_size=4,
stride=2,
norm_type=norm_type,
act_type=act_type,
mode=mode)
# 64, 128
conv4 = B.conv_block(base_nf * 2,
base_nf * 4,
kernel_size=3,
stride=1,
norm_type=norm_type,
act_type=act_type,
mode=mode)
conv5 = B.conv_block(base_nf * 4,
base_nf * 4,
kernel_size=4,
stride=2,
norm_type=norm_type,
act_type=act_type,
mode=mode)
self.feature_shallow = B.sequential(conv0, conv1, conv2, conv3, conv4,
conv5)
# 32, 256
conv6 = B.conv_block(base_nf * 4,
base_nf * 8,
kernel_size=3,
stride=1,
norm_type=norm_type,
act_type=act_type,
mode=mode)
conv7 = B.conv_block(base_nf * 8,
base_nf * 8,
kernel_size=4,
stride=2,
norm_type=norm_type,
act_type=act_type,
mode=mode)
# 16, 512
conv8 = B.conv_block(base_nf * 8,
base_nf * 8,
kernel_size=3,
stride=1,
norm_type=norm_type,
act_type=act_type,
mode=mode)
conv9 = B.conv_block(base_nf * 8,
base_nf * 8,
kernel_size=4,
stride=2,
norm_type=norm_type,
act_type=act_type,
mode=mode)
# 8, 512
conv10 = B.conv_block(base_nf * 8,
base_nf * 8,
kernel_size=3,
stride=1,
norm_type=norm_type,
act_type=act_type,
mode=mode)
conv11 = B.conv_block(base_nf * 8,
base_nf * 8,
kernel_size=4,
stride=2,
norm_type=norm_type,
act_type=act_type,
mode=mode)
# 4, 512
self.feature_deep = B.sequential(conv6, conv7, conv8, conv9, conv10,
conv11)
# classifier
self.classifier = nn.Sequential(nn.Linear(512 * 4 * 4, 100),
nn.LeakyReLU(0.2, True),
nn.Linear(100, 1))
def __init__(self,
in_nc=3,
out_nc=3,
upscale=4,
nf=64,
res_type='res',
n_mid=2,
n_tail=2,
n_HG=6,
act_type='leakyrelu',
inter_supervis=True,
mscale_inter_super=False,
share_upsample=False):
super(HourGlassNetMultiScaleInt, self).__init__()
self.n_HG = n_HG
self.inter_supervis = inter_supervis
if upscale == 3:
ksize = 3
else:
ksize = 1
self.conv_in = B.conv_block(in_nc,
nf,
kernel_size=3,
norm_type=None,
act_type=None)
def make_upsample_block(upscale=4, in_ch=64, out_nc=3, kernel_size=3):
n_upscale = 1 if upscale == 3 else int(math.log(upscale, 2))
LR_conv = B.conv_block(nf,
nf,
kernel_size=3,
norm_type=None,
act_type=None,
mode='CNA')
HR_conv0 = B.conv_block(nf,
nf,
kernel_size=3,
norm_type=None,
act_type='leakyrelu')
HR_conv1 = B.conv_block(nf,
out_nc,
kernel_size=kernel_size,
norm_type=None,
act_type=None)
if upscale == 1:
return nn.Sequential(LR_conv, HR_conv0, HR_conv1)
elif upscale == 3:
upsampler = B.upconv_blcok(nf, nf, 3, act_type=act_type)
else:
upsampler = [
B.upconv_blcok(nf, nf, act_type=act_type)
for _ in range(n_upscale)
]
return nn.Sequential(LR_conv, *upsampler, HR_conv0, HR_conv1)
#Actually, the ksize can be any size for all scales
self.flat_map = make_upsample_block(upscale=upscale, kernel_size=ksize)
self.edge_map = make_upsample_block(upscale=upscale, kernel_size=ksize)
self.corner_map = make_upsample_block(upscale=upscale,
kernel_size=ksize)
self.upsample_flat = make_upsample_block(upscale=upscale)
self.upsample_edge = make_upsample_block(upscale=upscale)
self.upsample_corner = make_upsample_block(upscale=upscale)
for i in range(n_HG):
if i != n_HG - 1:
HG_block = HourGlassBlock(res_type=res_type,
n_mid=n_mid,
n_tail=n_tail)
else:
HG_block = HourGlassBlock(res_type=res_type,
n_mid=n_mid,
n_tail=0)
setattr(self, 'HG_%d' % i, HG_block)
def __init__(self,
lm_nc,
pan_nc=1,
base_nf=64,
downscale=4,
norm_type='batch',
act_type='leakyrelu',
mode='CNA'):
super().__init__()
assert downscale in (2, 4), 'downscale should be 2 or 4'
# LM
# input (bs, c, 64, 64)
# output (bs, 32, 64, 64)
conv_LM0 = B.conv_block(lm_nc,
base_nf // 2,
kernel_size=3,
norm_type=None,
act_type=act_type,
mode=mode)
conv_LM1 = B.conv_block(base_nf // 2,
base_nf // 2,
kernel_size=3,
stride=1,
norm_type=norm_type,
act_type=act_type,
mode=mode)
self.LM = B.sequential(conv_LM0, conv_LM1)
# PAN
# input (bs, 1, 256, 256)
# output (bs, 32, 64, 64)
conv_PAN0 = B.conv_block(pan_nc,
base_nf // 4,
kernel_size=3,
norm_type=None,
act_type=act_type,
mode=mode)
# size down -> 128
conv_PAN1 = B.conv_block(base_nf // 4,
base_nf // 4,
kernel_size=4,
stride=2,
norm_type=norm_type,
act_type=act_type,
mode=mode) if downscale == 4 else None
conv_PAN2 = B.conv_block(base_nf // 4,
base_nf // 2,
kernel_size=3,
stride=1,
norm_type=norm_type,
act_type=act_type,
mode=mode) if downscale == 4 else None
# size down -> 64
conv_PAN3 = B.conv_block(base_nf // 2,
base_nf // 2,
kernel_size=4,
stride=2,
norm_type=norm_type,
act_type=act_type,
mode=mode)
self.PAN = B.sequential(conv_PAN0, conv_PAN1, conv_PAN2, conv_PAN3)
# Fus(Concatenate(LM_output, PAN_output))
# input (bs, 64, 64, 64)
# output (bs, 64, 64, 64)
conv_Fus = B.conv_block(base_nf,
base_nf,
kernel_size=3,
norm_type=None,
act_type=act_type,
mode=mode)
self.Fus = B.sequential(conv_Fus)
# head
# input (bs, 64, 64, 64)
# output (bs, 512, 4, 4)
convs = []
this_nf = base_nf
while True:
if this_nf < 512:
nf123 = this_nf
nf4 = 2 * this_nf
# elif this_nf == 512:
# nf123 = 512
# nf4 = 512
else:
break
# size down
size_down = B.conv_block(nf123,
nf123,
kernel_size=4,
stride=2,
norm_type=norm_type,
act_type=act_type,
mode=mode)
# channels up
size_keep = B.conv_block(nf123,
nf4,
kernel_size=3,
stride=1,
norm_type=norm_type,
act_type=act_type,
mode=mode)
convs.extend([size_down, size_keep])
this_nf *= 2
# 8, 512
conv_last = B.conv_block(512,
512,
kernel_size=4,
stride=2,
norm_type=norm_type,
act_type=act_type,
mode=mode)
# 4, 512
convs.append(conv_last)
self.head = B.sequential(*convs)
# classifier
self.classifier = nn.Sequential(nn.Linear(512 * 4 * 4, 100),
nn.LeakyReLU(0.2, True),
nn.Linear(100, 1))
def get_symbol(num_classes):
m = config.multiplier
sFilters = max(int(64*m), 32)
mFilters = max(int(128*m), 32)
nFilters = int(256*m)
nModules = 1
bn_mom = config.bn_mom
workspace = config.workspace
nStacks = config.net_stacks
binarize = config.net_binarize
input_size = config.input_img_size
label_size = config.output_label_size
use_coherent = config.net_coherent
use_SAT = config.net_sat
N = config.net_n
DCN = config.net_dcn
per_batch_size = config.per_batch_size
print('binarize', binarize)
print('use_coherent', use_coherent)
print('use_SAT', use_SAT)
print('use_N', N)
print('use_DCN', DCN)
print('per_batch_size', per_batch_size)
#assert(label_size==64 or label_size==32)
#assert(input_size==128 or input_size==256)
coherentor = SymCoherent(per_batch_size)
D = input_size // label_size
print(input_size, label_size, D)
data = mx.sym.Variable(name='data')
data = data-127.5
data = data*0.0078125
gt_label = mx.symbol.Variable(name='softmax_label')
losses = []
closses = []
ref_label = gt_label
if D==4:
body = Conv(data=data, num_filter=sFilters, kernel=(7, 7), stride=(2,2), pad=(3, 3),
no_bias=True, name="conv0", workspace=workspace)
else:
body = Conv(data=data, num_filter=sFilters, kernel=(3, 3), stride=(1,1), pad=(1, 1),
no_bias=True, name="conv0", workspace=workspace)
body = mx.sym.BatchNorm(data=body, fix_gamma=False, eps=2e-5, momentum=bn_mom, name='bn0')
body = Act(data=body, act_type='relu', name='relu0')
dcn = False
body = conv_block(body, mFilters, (1,1), sFilters==mFilters, 'res0', False, dcn, 1)
body = mx.sym.Pooling(data=body, kernel=(2, 2), stride=(2,2), pad=(0,0), pool_type='max')
body = conv_block(body, mFilters, (1,1), True, 'res1', False, dcn, 1) #TODO
body = conv_block(body, nFilters, (1,1), mFilters==nFilters, 'res2', binarize, dcn, 1) #binarize=True?
heatmap = None
for i in xrange(nStacks):
shortcut = body
body = hourglass(body, nFilters, nModules, config.net_n, workspace, 'stack%d_hg'%(i), binarize, dcn)
for j in xrange(nModules):
body = conv_block(body, nFilters, (1,1), True, 'stack%d_unit%d'%(i,j), binarize, dcn, 1)
_dcn = True if config.net_dcn>=2 else False
ll = ConvFactory(body, nFilters, (1,1), dcn = _dcn, name='stack%d_ll'%(i))
_name = "heatmap%d"%(i) if i<nStacks-1 else "heatmap"
_dcn = True if config.net_dcn>=2 else False
if not _dcn:
out = Conv(data=ll, num_filter=num_classes, kernel=(1, 1), stride=(1,1), pad=(0,0),
name=_name, workspace=workspace)
else:
out_offset = mx.symbol.Convolution(name=_name+'_offset', data = ll,
num_filter=18, pad=(1, 1), kernel=(3, 3), stride=(1, 1))
out = mx.contrib.symbol.DeformableConvolution(name=_name, data=ll, offset=out_offset,
num_filter=num_classes, pad=(1,1), kernel=(3, 3), num_deformable_group=1, stride=(1, 1), dilate=(1, 1), no_bias=False)
#out = Conv(data=ll, num_filter=num_classes, kernel=(3,3), stride=(1,1), pad=(1,1),
# name=_name, workspace=workspace)
if i<nStacks-1:
ll2 = Conv(data=ll, num_filter=nFilters, kernel=(1, 1), stride=(1,1), pad=(0,0),
name="stack%d_ll2"%(i), workspace=workspace)
out2 = Conv(data=out, num_filter=nFilters, kernel=(1, 1), stride=(1,1), pad=(0,0),
name="stack%d_out2"%(i), workspace=workspace)
body = mx.symbol.add_n(shortcut, ll2, out2)
_dcn = True if (config.net_dcn==1 or config.net_dcn==3) else False
if _dcn:
_name = "stack%d_out3" % (i)
out3_offset = mx.symbol.Convolution(name=_name+'_offset', data = body,
num_filter=18, pad=(1, 1), kernel=(3, 3), stride=(1, 1))
out3 = mx.contrib.symbol.DeformableConvolution(name=_name, data=body, offset=out3_offset,
num_filter=nFilters, pad=(1,1), kernel=(3, 3), num_deformable_group=1, stride=(1, 1), dilate=(1, 1), no_bias=False)
body = out3
elif i==nStacks-1:
heatmap = out
# loss = ce_loss(out, ref_label)
# loss = loss/nStacks
loss = l2_loss(out, ref_label)
losses.append(loss)
if config.net_coherent>0:
ux, dx = coherentor.get(out)
closs = l2_loss(ux, dx)
closs = closs/nStacks
closses.append(closs)
pred = mx.symbol.BlockGrad(heatmap)
#loss = mx.symbol.add_n(*losses)
#loss = mx.symbol.MakeLoss(loss)
#syms = [loss]
syms = []
for loss in losses:
loss = mx.symbol.MakeLoss(loss)
syms.append(loss)
if len(closses)>0:
coherent_weight = 0.0001
closs = mx.symbol.add_n(*closses)
closs = mx.symbol.MakeLoss(closs, grad_scale = coherent_weight)
syms.append(closs)
syms.append(pred)
sym = mx.symbol.Group( syms )
return sym
def __init__(self,
in_nc,
base_nf,
norm_type='batch',
act_type='leakyrelu',
mode='CNA'):
super(Discriminator_VGG_128, self).__init__()
# features
# h and w, c
# input 128, 3
conv0 = B.conv_block(in_nc,
base_nf,
kernel_size=3,
norm_type=None,
act_type=act_type,
mode=mode)
conv1 = B.conv_block(base_nf,
base_nf,
kernel_size=4,
stride=2,
norm_type=norm_type,
act_type=act_type,
mode=mode)
# 64, 64
conv2 = B.conv_block(base_nf,
base_nf * 2,
kernel_size=3,
stride=1,
norm_type=norm_type,
act_type=act_type,
mode=mode)
conv3 = B.conv_block(base_nf * 2,
base_nf * 2,
kernel_size=4,
stride=2,
norm_type=norm_type,
act_type=act_type,
mode=mode)
# 32, 128
conv4 = B.conv_block(base_nf * 2,
base_nf * 4,
kernel_size=3,
stride=1,
norm_type=norm_type,
act_type=act_type,
mode=mode)
conv5 = B.conv_block(base_nf * 4,
base_nf * 4,
kernel_size=4,
stride=2,
norm_type=norm_type,
act_type=act_type,
mode=mode)
# 16, 256
conv6 = B.conv_block(base_nf * 4,
base_nf * 8,
kernel_size=3,
stride=1,
norm_type=norm_type,
act_type=act_type,
mode=mode)
conv7 = B.conv_block(base_nf * 8,
base_nf * 8,
kernel_size=4,
stride=2,
norm_type=norm_type,
act_type=act_type,
mode=mode)
# 8, 512
conv8 = B.conv_block(base_nf * 8,
base_nf * 8,
kernel_size=3,
stride=1,
norm_type=norm_type,
act_type=act_type,
mode=mode)
conv9 = B.conv_block(base_nf * 8,
base_nf * 8,
kernel_size=4,
stride=2,
norm_type=norm_type,
act_type=act_type,
mode=mode)
# 4, 512
self.features = B.sequential(conv0, conv1, conv2, conv3, conv4, conv5,
conv6, conv7, conv8, conv9)
# classifier
self.classifier = nn.Sequential(nn.Linear(512 * 4 * 4, 100),
nn.LeakyReLU(0.2, True),
nn.Linear(100, 1))
