def __init__(self,
in_nc=4,
out_nc=3,
nc=[64, 128, 256, 512],
nb=2,
act_mode='R',
downsample_mode='strideconv',
upsample_mode='convtranspose'):
super(ResUNet_decoder, self).__init__()
# upsample
if upsample_mode == 'upconv':
upsample_block = B.upsample_upconv
elif upsample_mode == 'pixelshuffle':
upsample_block = B.upsample_pixelshuffle
elif upsample_mode == 'convtranspose':
upsample_block = B.upsample_convtranspose
else:
raise NotImplementedError(
'upsample mode [{:s}] is not found'.format(upsample_mode))
self.m_up3 = B.sequential(
upsample_block(nc[3], nc[2], bias=False, mode='2'), *[
B.ResBlock(nc[2], nc[2], bias=False, mode='C' + act_mode + 'C')
for _ in range(nb)
])
self.m_up2 = B.sequential(
upsample_block(nc[2], nc[1], bias=False, mode='2'), *[
B.ResBlock(nc[1], nc[1], bias=False, mode='C' + act_mode + 'C')
for _ in range(nb)
])
self.m_up1 = B.sequential(
upsample_block(nc[1], nc[0], bias=False, mode='2'), *[
B.ResBlock(nc[0], nc[0], bias=False, mode='C' + act_mode + 'C')
for _ in range(nb)
])
# self.m_tail = B.conv(nc[0], out_nc,kernel_size=7,padding=3, bias=False, mode='C')
self.m_tail = B.conv(nc[0], out_nc, bias=False, mode='C')
self.m_final1 = B.conv(out_nc,
out_nc,
kernel_size=1,
bias=False,
mode='C')
self.m_final2 = B.conv(out_nc,
out_nc,
kernel_size=1,
bias=False,
mode='C')
def __init__(self,
in_nc=4,
out_nc=3,
nc=[64, 128, 256, 512],
nb=2,
act_mode='R',
downsample_mode='strideconv',
upsample_mode='convtranspose'):
super(ResUNet_encoder, self).__init__()
# self.m_head = B.conv(in_nc, nc[0],kernel_size=7,padding=3, bias=False, mode='C')
self.m_head = B.conv(in_nc, nc[0], bias=False, mode='C')
# downsample
if downsample_mode == 'avgpool':
downsample_block = B.downsample_avgpool
elif downsample_mode == 'maxpool':
downsample_block = B.downsample_maxpool
elif downsample_mode == 'strideconv':
downsample_block = B.downsample_strideconv
else:
raise NotImplementedError(
'downsample mode [{:s}] is not found'.format(downsample_mode))
self.m_down1 = B.sequential(
*[
B.ResBlock(nc[0], nc[0], bias=False, mode='C' + act_mode + 'C')
for _ in range(nb)
], downsample_block(nc[0], nc[1], bias=False, mode='2'))
self.m_down2 = B.sequential(
*[
B.ResBlock(nc[1], nc[1], bias=False, mode='C' + act_mode + 'C')
for _ in range(nb)
], downsample_block(nc[1], nc[2], bias=False, mode='2'))
self.m_down3 = B.sequential(
*[
B.ResBlock(nc[2], nc[2], bias=False, mode='C' + act_mode + 'C')
for _ in range(nb)
], downsample_block(nc[2], nc[3], bias=False, mode='2'))
self.m_body = B.sequential(*[
B.ResBlock(nc[3], nc[3], bias=False, mode='C' + act_mode + 'C')
for _ in range(nb)
])