def get_netG():
"""Get net G"""
# build the generator
netG = nn.Sequential()
with netG.name_scope():
# input is Z, going into a convolution
netG.add(nn.Conv2DTranspose(ngf * 8, 4, 1, 0, use_bias=False))
netG.add(nn.BatchNorm())
netG.add(nn.Activation('relu'))
# state size. (ngf*8) x 4 x 4
netG.add(nn.Conv2DTranspose(ngf * 4, 4, 2, 1, use_bias=False))
netG.add(nn.BatchNorm())
netG.add(nn.Activation('relu'))
# state size. (ngf*4) x 8 x 8
netG.add(nn.Conv2DTranspose(ngf * 2, 4, 2, 1, use_bias=False))
netG.add(nn.BatchNorm())
netG.add(nn.Activation('relu'))
# state size. (ngf*2) x 16 x 16
netG.add(nn.Conv2DTranspose(ngf, 4, 2, 1, use_bias=False))
netG.add(nn.BatchNorm())
netG.add(nn.Activation('relu'))
# state size. (ngf) x 32 x 32
netG.add(nn.Conv2DTranspose(nc, 4, 2, 1, use_bias=False))
netG.add(nn.Activation('tanh'))
# state size. (nc) x 64 x 64
return netG
def __init__(self, dim, **kwargs):
super(Generator, self).__init__(**kwargs)
self.dim = dim
self.preprocess = nn.HybridSequential()
with self.preprocess.name_scope():
self.preprocess.add(
nn.Conv2D(4 * 4 * 4 * dim, kernel_size=1, use_bias=False))
self.preprocess.add(nn.BatchNorm())
self.preprocess.add(nn.Activation('relu'))
self.block1 = nn.HybridSequential()
with self.block1.name_scope():
self.block1.add(
nn.Conv2DTranspose(2 * dim,
kernel_size=2,
strides=2,
use_bias=False))
self.block1.add(nn.BatchNorm())
self.block1.add((nn.Activation('relu')))
self.block2 = nn.HybridSequential()
with self.block2.name_scope():
self.block2.add(
nn.Conv2DTranspose(dim,
kernel_size=2,
strides=2,
use_bias=False))
self.block2.add(nn.BatchNorm())
self.block2.add((nn.Activation('relu')))
self.deconv_out = nn.Conv2DTranspose(3, kernel_size=2, strides=2)
def __init__(self, fwblock, name, **kwargs):
super(topdown_residual_unit, self).__init__(**kwargs)
self.dim_match = fwblock.dim_match
self.num_filter = fwblock.num_filter
self.out_channels = fwblock.in_channels
self.ratio = fwblock.ratio
self.strides = fwblock.strides
self.num_group = fwblock.num_group
self.bn_mom = fwblock.bn_mom
output_padding = 1 if self.strides[0] > 1 else 0
if not self.dim_match:
self.fc_sc = nn.Conv2DTranspose(channels=self.out_channels, in_channels=self.num_filter, kernel_size=(1,1), strides=self.strides, use_bias=False, output_padding=output_padding, params=fwblock.fc_sc.params, prefix=name + '_tdsc_dense_')
# block 3
self.bn3 = nn.BatchNorm(in_channels=self.num_filter, epsilon=2e-5, momentum=self.bn_mom, prefix=name + '_td_batchnorm3_')
self.bn3min = nn.BatchNorm(in_channels=self.num_filter, epsilon=2e-5, momentum=self.bn_mom, prefix=name + '_td_batchnorm3min_')
self.conv3 = nn.Conv2DTranspose(channels=int(self.num_filter*0.5), in_channels=self.num_filter, kernel_size=(1,1), use_bias=False, params=fwblock.conv3.params, prefix=name + '_td_conv3_')
# block 2
self.bn2 = nn.BatchNorm(in_channels=int(self.num_filter*0.5), epsilon=2e-5, momentum=self.bn_mom, prefix=name + '_td_batchnorm2_')
self.bn2min = nn.BatchNorm(in_channels=int(self.num_filter*0.5), epsilon=2e-5, momentum=self.bn_mom, prefix=name + '_td_batchnorm2min_')
self.conv2 = nn.Conv2DTranspose(channels=int(self.num_filter*0.5), in_channels=int(self.num_filter*0.5), kernel_size=(3,3), strides=self.strides, padding=(1,1), output_padding=output_padding, use_bias=False, groups=self.num_group, params=fwblock.conv2.params, prefix=name + '_td_conv2_')
# block 1
self.bn1 = nn.BatchNorm(in_channels=int(self.num_filter*0.5), epsilon=2e-5, momentum=self.bn_mom, prefix=name + '_td_batchnorm1_')
self.bn1min = nn.BatchNorm(in_channels=int(self.num_filter*0.5), epsilon=2e-5, momentum=self.bn_mom, prefix=name + '_td_batchnorm1min_')
self.conv1 = nn.Conv2DTranspose(channels=self.out_channels, in_channels=int(self.num_filter*0.5), kernel_size=(1,1), use_bias=False, params=fwblock.conv1.params, prefix=name + '_td_conv1_')
def __init__(self, batch_images, classes, mask_channels, num_fcn_convs=0, norm_layer=None,
norm_kwargs=None, **kwargs):
super(Mask, self).__init__(**kwargs)
self._batch_images = batch_images
self.classes = classes
init = mx.init.Xavier(rnd_type='gaussian', factor_type='out', magnitude=2)
with self.name_scope():
if num_fcn_convs > 0:
self.deconv = nn.HybridSequential()
for _ in range(num_fcn_convs):
self.deconv.add(
nn.Conv2D(mask_channels, kernel_size=(3, 3), strides=(1, 1),
padding=(1, 1), weight_initializer=init))
if norm_layer is not None and norm_layer is SyncBatchNorm:
self.deconv.add(norm_layer(**norm_kwargs))
self.deconv.add(nn.Activation('relu'))
self.deconv.add(
nn.Conv2DTranspose(mask_channels, kernel_size=(2, 2), strides=(2, 2),
padding=(0, 0), weight_initializer=init))
if norm_layer is not None and norm_layer is SyncBatchNorm:
self.deconv.add(norm_layer(**norm_kwargs))
else:
# this is for compatibility of older models.
self.deconv = nn.Conv2DTranspose(mask_channels, kernel_size=(2, 2), strides=(2, 2),
padding=(0, 0), weight_initializer=init)
self.mask = nn.Conv2D(len(classes), kernel_size=(1, 1), strides=(1, 1), padding=(0, 0),
weight_initializer=init)
def getDecoder(self):
decoder = nn.HybridSequential(prefix='decoder')
decoder.add(
nn.Dense(units=1024)) # 1024 for size=64, 9216 for size=128
decoder.add(Expand_dims(axis=-1))
decoder.add(Expand_dims(axis=-1))
# relu deconv 128x5 to 5x5x128
decoder.add(
nn.Conv2DTranspose(channels=128,
kernel_size=5,
strides=2,
activation='relu'))
# relu deconv 64x5 to 13x13x64
decoder.add(
nn.Conv2DTranspose(channels=64,
kernel_size=5,
strides=2,
activation='relu'))
# relu deconv 32x6 to 30x30x32
# decoder.add(nn.Conv2DTranspose(channels=32, kernel_size=6, strides=2, activation='relu'))
decoder.add(
nn.Conv2DTranspose(channels=32,
kernel_size=6,
strides=2,
activation='relu'))
# decoder.add(nn.Conv2DTranspose(channels=16, kernel_size=6, strides=2, activation='relu'))
# sigmoid deconv 3x6 to 64x64x3
decoder.add(
nn.Conv2DTranspose(channels=3,
kernel_size=6,
strides=2,
activation='sigmoid'))
# decoder.add(Reshape((3, 128, 128)))
return decoder
def __init__(self, isize, nz, nc, ngf, ngpu, n_extra_layers=0):
super(DCGAN_G_nobn, self).__init__()
self.ngpu = ngpu
assert isize % 16 == 0, "isize has to be a multiple of 16"
cngf, tisize = ngf // 2, 4
while tisize != isize:
cngf = cngf * 2
tisize = tisize * 2
with self.name_scope():
main = nn.HybridSequential()
main.add(
nn.Conv2DTranspose(in_channels=nz, channels=cngf, kernel_size=4, strides=1, padding=0, use_bias=False,
activation='relu', prefix='initial.{0}-{1}.convt'.format(nz, cngf)))
csize, cndf = 4, cngf
while csize < isize // 2:
main.add(nn.Conv2DTranspose(in_channels=cngf, channels=cngf // 2, kernel_size=4, strides=2, padding=1,
use_bias=False, activation='relu',
prefix='pyramid.{0}-{1}.convt'.format(cngf, cngf // 2)))
cngf = cngf // 2
csize = csize * 2
# Extra layers
for t in range(n_extra_layers):
main.add(nn.Conv2D(in_channels=cngf, channels=cngf, kernel_size=3, strides=1, padding=1, use_bias=False,
activation='relu', prefix='extra-layers-{0}.{1}.conv'.format(t, cngf)))
main.add(
nn.Conv2DTranspose(in_channels=cngf, channels=nc, kernel_size=4, strides=2, padding=1, use_bias=False,
activation='tanh', prefix='final.{0}-{1}.convt'.format(cngf, nc)))
self.main = main
def get_generator_block(self,
input_dim,
output_dim,
kernel_size=3,
strides=2,
final_layer=False):
layer = nn.HybridSequential()
if not final_layer:
layer.add(
nn.Conv2DTranspose(in_channels=input_dim,
channels=output_dim,
kernel_size=kernel_size,
strides=strides,
use_bias=False),
nn.BatchNorm(in_channels=output_dim),
nn.Activation(activation="relu"))
else:
layer.add(
nn.Conv2DTranspose(in_channels=input_dim,
channels=output_dim,
kernel_size=kernel_size,
strides=strides,
use_bias=False),
nn.Activation(activation="tanh"))
return layer
def __init__(self, **kwargs):
super(Generator, self).__init__(**kwargs)
with self.name_scope():
self.generator = nn.HybridSequential(prefix='generator')
# input is Z, going into a convolution
self.generator.add(
nn.Conv2DTranspose(64 * 8, 4, 1, 0, use_bias=False))
self.generator.add(nn.BatchNorm())
self.generator.add(nn.Activation('relu'))
# state size (64*8)x4x4
self.generator.add(
nn.Conv2DTranspose(64 * 4, 4, 2, 1, use_bias=False))
self.generator.add(nn.BatchNorm())
self.generator.add(nn.Activation('relu'))
# state size (64*4)x8x8
self.generator.add(
nn.Conv2DTranspose(64 * 2, 4, 2, 1, use_bias=False))
self.generator.add(nn.BatchNorm())
self.generator.add(nn.Activation('relu'))
# state size (64*2)x16x16
self.generator.add(
nn.Conv2DTranspose(64 * 1, 4, 2, 1, use_bias=False))
self.generator.add(nn.BatchNorm())
self.generator.add(nn.Activation('relu'))
# state size (64*1)x32x32
self.generator.add(nn.Conv2DTranspose(3, 4, 2, 1, use_bias=False))
self.generator.add(nn.Activation('tanh'))
def create_model(num_channels=3, ngf=64, ndf=64):
netG = nn.Sequential()
with netG.name_scope():
# input shape: (?, 1000 + random_input_size, 1, 1)
netG.add(
nn.Conv2DTranspose(channels=ngf * 8,
kernel_size=4,
strides=1,
padding=0,
use_bias=False))
netG.add(nn.BatchNorm())
netG.add(nn.Activation('relu'))
# state shape: (?, ngf * 8, 4, 4)
netG.add(
nn.Conv2DTranspose(channels=ngf * 4,
kernel_size=4,
strides=2,
padding=1,
use_bias=False))
netG.add(nn.BatchNorm())
netG.add(nn.Activation('relu'))
# state shape: (?, ngf * 4, 8, 8)
netG.add(
nn.Conv2DTranspose(channels=ngf * 2,
kernel_size=4,
strides=2,
padding=1,
use_bias=False))
netG.add(nn.BatchNorm())
netG.add(nn.Activation('relu'))
# state shape: (?, ngf * 2, 16, 16)
netG.add(
nn.Conv2DTranspose(channels=ngf,
kernel_size=4,
strides=2,
padding=1,
use_bias=False))
netG.add(nn.BatchNorm())
netG.add(nn.Activation('relu'))
# state shape: (?, ngf * 2, 32, 32)
netG.add(
nn.Conv2DTranspose(channels=num_channels,
kernel_size=4,
strides=2,
padding=1,
use_bias=False))
netG.add(nn.Activation('tanh'))
# state shape: (?, num_channels, 64, 64)
netD = Discriminator(ndf)
return netG, netD
def __init__(self, features, config, **kwargs):
super().__init__(**kwargs)
self.feature_layers = [4, 5, 6, 7]
self.feature_dims = [64, 128, 256, 512]
self.features = features
with self.name_scope():
# self.backbone = vision.resnet50_v1()
self.add_layers('reduce_dim', conv_bn_relu(64, 1),
conv_bn_relu(64, 1), conv_bn_relu(64, 1),
conv_bn_relu(64, 1))
self.add_layers(
'deconv',
nn.Conv2DTranspose(
64,
4,
strides=2,
padding=1,
weight_initializer=mx.initializer.MSRAPrelu(slope=0.1)),
nn.Conv2DTranspose(
64,
4,
strides=2,
padding=1,
weight_initializer=mx.initializer.MSRAPrelu(slope=0.1)),
nn.Conv2DTranspose(
64,
4,
strides=2,
padding=1,
weight_initializer=mx.initializer.MSRAPrelu(slope=0.1)))
self.add_layers('pred', self._prediction_head(),
self._prediction_head(), self._prediction_head(),
self._prediction_head())
self.flow = nn.HybridSequential(prefix='flow')
self.flow.add(
nn.Conv2D(64,
7,
padding=3,
weight_initializer=Xavier(rnd_type='gaussian',
magnitude=2)))
self.flow.add(nn.LeakyReLU(0.1))
self.flow.add(
nn.Conv2D(64,
7,
padding=3,
weight_initializer=Xavier(rnd_type='gaussian',
magnitude=2)))
self.flow.add(nn.LeakyReLU(0.1))
self.flow.add(
nn.Conv2D(64,
7,
padding=3,
weight_initializer=Xavier(rnd_type='gaussian',
magnitude=2)))
self.flow.add(nn.LeakyReLU(0.1))
def __init__(self,
n_latent=512,
n_channels=3,
out_width=64,
out_height=64,
n_base_channels=16):
super(ConvDecoder, self).__init__()
# Store some of the hyper paramteres
self.n_latent = n_latent
self.n_channels = n_channels
self.out_width = out_width
self.out_height = out_height
self.n_base_channels = n_base_channels
# Construct the decoder network
with self.name_scope():
# Construct the decoder network
# For the decoder network, its input is 4-dimensional arrays
# of shape (batch_size, n_latent, 1, 1)
# Decoder's architecture came from Deep Convolutional
# Generative Adversarial Network tutorial from MXNet
self.decoder = nn.Sequential(prefix='decoder')
# Add convolution layers with decreasing number of channels
self.decoder.add(
nn.Conv2DTranspose(n_base_channels * 8,
4,
1,
0,
use_bias=False), nn.BatchNorm(),
nn.Activation('relu'))
self.decoder.add(
nn.Conv2DTranspose(n_base_channels * 4,
4,
2,
1,
use_bias=False), nn.BatchNorm(),
nn.Activation('relu'))
self.decoder.add(
nn.Conv2DTranspose(n_base_channels * 2,
4,
2,
1,
use_bias=False), nn.BatchNorm(),
nn.Activation('relu'))
self.decoder.add(
nn.Conv2DTranspose(n_base_channels * 1,
4,
2,
1,
use_bias=False), nn.BatchNorm(),
nn.Activation('relu'))
# The last layer uses sigmoid because I want image pixel data between 0 and 1
# if it doesn't work I will use tanh as the tutorial did
self.decoder.add(
nn.Conv2DTranspose(self.n_channels, 4, 2, 1, use_bias=False),
nn.Activation('sigmoid'))
def __init__(self, block_unit, out_channels, channels, cardinality,
bottleneck_width, stride, **kwargs):
super(TopdownBlock, self).__init__(**kwargs)
D = int(math.floor(channels * (bottleneck_width / 64)))
group_width = cardinality * D
output_padding = 1 if stride > 1 else 0
if block_unit.downsample is not None:
self.upsample = nn.HybridSequential(prefix='')
self.upsample.add(nn.BatchNorm())
self.upsample.add(
nn.Conv2DTranspose(channels=out_channels,
in_channels=channels * 4,
kernel_size=1,
strides=stride,
use_bias=False,
output_padding=output_padding,
params=block_unit.downsample[0].params))
else:
self.upsample = None
self.body = nn.HybridSequential(prefix='')
self.body.add(nn.BatchNorm())
self.body.add(
nn.Conv2DTranspose(channels=group_width,
in_channels=channels * 4,
kernel_size=1,
use_bias=False,
params=block_unit.body[-2].params))
self.body.add(nn.BatchNorm())
self.body.add(
nn.Conv2DTranspose(channels=group_width,
in_channels=group_width,
kernel_size=3,
strides=stride,
padding=1,
output_padding=output_padding,
use_bias=False,
groups=cardinality,
params=block_unit.body[-5].params))
self.body.add(nn.BatchNorm())
self.body.add(
nn.Conv2DTranspose(channels=out_channels,
in_channels=group_width,
kernel_size=1,
use_bias=False,
params=block_unit.body[-8].params))
self.insnorms = nn.HybridSequential(prefix='insnorms_')
for layer in self.body:
if layer.name.find('batchnorm') != -1:
self.insnorms.add(InstanceNorm())
if self.upsample is not None:
self.insnormsds = nn.HybridSequential(prefix='insnormsds_')
for layer in self.upsample:
if layer.name.find('batchnorm') != -1:
self.insnormsds.add(InstanceNorm())
def __init__(self, size_img, num_z, num_hidden, num_c, num_extra_layers=0):
super(DCGAN_G, self).__init__()
cngf, t_size_img = num_hidden // 2, 4
while t_size_img != size_img:
cngf = cngf * 2
t_size_img = t_size_img * 2
num_hidden = cngf
with self.name_scope():
self.base = nn.Sequential()
# inpurt is Z (nz dim vector), mapping to a convNet
self.base.add(
nn.Conv2DTranspose(channels=num_hidden,
in_channels=num_z,
kernel_size=4,
strides=1,
padding=0,
use_bias=False))
self.base.add(nn.BatchNorm(in_channels=num_hidden))
self.base.add(nn.LeakyReLU(0))
size_conv = 4
# 不断 DeConv，直到 feature map 到达 size_img 的一半大小
while size_conv < size_img // 2:
self.base.add(
nn.Conv2DTranspose(channels=num_hidden // 2,
in_channels=num_hidden,
kernel_size=4,
strides=2,
padding=1,
use_bias=False))
self.base.add(nn.BatchNorm(in_channels=num_hidden // 2))
self.base.add(nn.LeakyReLU(0))
num_hidden = num_hidden // 2
size_conv *= 2
for _ in range(num_extra_layers):
self.base.add(
nn.Conv2D(channels=num_hidden,
in_channels=num_hidden,
kernel_size=3,
strides=1,
padding=1,
use_bias=False))
self.base.add(nn.BatchNorm(in_channels=num_hidden))
self.base.add(nn.LeakyReLU(0))
self.base.add(
nn.Conv2DTranspose(channels=num_c,
in_channels=num_hidden,
kernel_size=4,
strides=2,
padding=1,
use_bias=False,
activation='tanh'))
def make_generic_gen(isize, nz, nc, ngf, n_extra_layers=0, use_bn=True):
assert isize % 16 == 0, 'isize has to be a multiple of 16'
gen = nn.HybridSequential(prefix='generator')
cngf, tisize = ngf // 2, 4
while tisize != isize:
cngf *= 2
tisize *= 2
# add input z vector
gen.add(
nn.Conv2DTranspose(in_channels=nz,
channels=cngf,
kernel_size=4,
strides=1,
padding=0,
use_bias=False,
prefix='tconv_{}->{}-'.format(nz, cngf)))
if use_bn:
gen.add(
nn.BatchNorm(in_channels=cngf,
prefix='bn_{}->{}-'.format(nz, cngf)))
gen.add(nn.LeakyReLU(0.05, prefix='lrelu{}->{}-'.format(nz, cngf)))
csize, cndf = 4, cngf
while csize < isize // 2:
gen.add(
nn.Conv2DTranspose(in_channels=cngf,
channels=cngf // 2,
kernel_size=4,
strides=2,
padding=1,
use_bias=False))
if use_bn:
gen.add(nn.BatchNorm(in_channels=cngf // 2))
gen.add(nn.LeakyReLU(0.05))
cngf //= 2
csize *= 2
for extra in range(n_extra_layers):
gen.add(
nn.Conv2D(in_channels=cngf,
channels=cngf,
kernel_size=3,
strides=1,
padding=1,
use_bias=False))
if use_bn:
nn.BatchNorm(in_channels=cngf),
gen.add(nn.LeakyReLU(0.05))
gen.add(
nn.Conv2DTranspose(in_channels=cngf,
channels=nc,
kernel_size=4,
strides=2,
padding=1,
use_bias=False), nn.Activation('tanh'))
return gen
def __init__(self, n_dims=128, **kwargs):
super(Generator, self).__init__(**kwargs)
with self.name_scope():
self.deconv_z = nn.Conv2DTranspose(n_dims * 2, 4, 1, 0)
self.deconv_label = nn.Conv2DTranspose(n_dims * 2, 4, 1, 0)
self.deconv2 = nn.Conv2DTranspose(n_dims * 2, 4, 1, 0)
self.deconv3 = nn.Conv2DTranspose(n_dims, 4, 2, 1)
self.deconv4 = nn.Conv2DTranspose(1, 4, 2, 1)
self.bn_z = nn.BatchNorm()
self.bn_label = nn.BatchNorm()
self.bn2 = nn.BatchNorm()
self.bn3 = nn.BatchNorm()
def __init__(self,
batch_images,
num_classes,
channels,
norm_layer=nn.InstanceNorm,
norm_kwargs={},
**kwargs):
super(DrivableMaps, self).__init__(**kwargs)
self._batch_images = batch_images
init = mx.init.Xavier(rnd_type='gaussian',
factor_type='out',
magnitude=2)
with self.name_scope():
self.map_head = nn.HybridSequential()
self.map_head.add(
nn.Conv2D(in_channels=1024,
channels=channels * 8,
kernel_size=3,
padding=1))
self.map_head.add(norm_layer(**norm_kwargs))
self.map_head.add(nn.Activation('relu'))
self.map_head.add(
nn.Conv2DTranspose(channels * 4,
kernel_size=4,
strides=2,
padding=1,
weight_initializer=init))
self.map_head.add(norm_layer(**norm_kwargs))
self.map_head.add(nn.Activation('relu'))
self.map_head.add(nn.Conv2D(channels * 2, kernel_size=3,
padding=1))
self.map_head.add(norm_layer(**norm_kwargs))
self.map_head.add(nn.Activation('relu'))
self.map_head.add(nn.Dropout(0.1))
self.map_head.add(
nn.Conv2DTranspose(channels,
kernel_size=4,
strides=2,
padding=1,
weight_initializer=init))
self.map_head.add(norm_layer(**norm_kwargs))
self.map_head.add(nn.Activation('relu'))
self.map_head.add(
nn.Conv2D(num_classes,
kernel_size=(1, 1),
strides=(1, 1),
padding=(0, 0),
weight_initializer=init))
def __init__(self,
outer_nc,
inner_nc,
submodule=None,
outermost=False,
innermost=False,
use_dropout=False):
super(UnetSkipConnectionBlock, self).__init__()
self.outermost = outermost
downconv = nn.Conv2D(inner_nc, kernel_size=4, strides=2, padding=1)
downrelu = nn.LeakyReLU(0.2)
downnorm = nn.InstanceNorm()
uprelu = nn.Activation('relu')
upnorm = nn.InstanceNorm()
self.model = nn.HybridSequential()
with self.model.name_scope():
if outermost:
self.model.add(downconv)
if submodule is not None:
self.model.add(submodule)
self.model.add(
uprelu,
nn.Conv2DTranspose(outer_nc,
kernel_size=4,
strides=2,
padding=1), nn.Activation('tanh'))
elif innermost:
self.model.add(
downrelu, downconv, uprelu,
nn.Conv2DTranspose(outer_nc,
kernel_size=4,
strides=2,
padding=1), upnorm)
else:
self.model.add(
downrelu,
downconv,
downnorm,
)
if submodule is not None:
self.model.add(submodule)
self.model.add(
uprelu,
nn.Conv2DTranspose(outer_nc,
kernel_size=4,
strides=2,
padding=1),
upnorm,
)
if use_dropout:
self.model.add(nn.Dropout(0.5))
def __init__(self, **kwargs):
super(Generator, self).__init__(**kwargs)
with self.name_scope():
self.dense_z = nn.Dense(256)
self.dense_label = nn.Dense(256)
self.deconv2 = nn.Conv2DTranspose(256, kernel_size=7)
self.deconv3 = nn.Conv2DTranspose(32, kernel_size=4, strides=2,padding=1)
self.deconv4 = nn.Conv2DTranspose(1, kernel_size=4, strides=2,padding=1)
self.bn_z = nn.BatchNorm()
self.bn_label = nn.BatchNorm()
self.bn2 = nn.BatchNorm()
self.bn3 = nn.BatchNorm()
self.bn4 = nn.BatchNorm()
def __init__(self, vgg_name, batch_size, num_class, use_bias=False, use_bn=False, do_topdown=False, do_countpath=False, do_pn=False, relu_td=False, do_nn=False):
super(VGG_DRM, self).__init__()
self.num_class = num_class
self.do_topdown = do_topdown
self.do_countpath = do_countpath
self.do_pn = do_pn
self.relu_td = relu_td
self.do_nn = do_nn
self.use_bn = use_bn
self.use_bias = use_bias
self.batch_size = batch_size
self.features, layers_drm, layers_drm_cp = self._make_layers(cfg[vgg_name], use_bias, use_bn, self.do_topdown, self.do_countpath)
with self.name_scope():
self.classifier = nn.HybridSequential(prefix='classifier_')
conv_layer = nn.Conv2D(in_channels=cfg[vgg_name][0][-2], channels=self.num_class, kernel_size=(1, 1), use_bias=True)
self.classifier.add(conv_layer)
self.classifier.add(nn.Flatten())
if self.do_topdown:
layers_drm += [nn.Conv2DTranspose(channels=cfg[vgg_name][0][-2], in_channels=self.num_class, kernel_size=(1,1), strides=(1, 1),
use_bias=False, params=conv_layer.params),
Reshape(shape=(self.num_class, 1, 1))]
with self.name_scope():
self.drm = nn.HybridSequential(prefix='drmtd_')
for block in layers_drm[::-1]:
self.drm.add(block)
if self.do_pn:
with self.name_scope():
self.insnorms = nn.HybridSequential(prefix='instancenorm_')
count_bn = 0
for i in range(len(self.drm._children)):
if (self.drm._children[i].name.find('batchnorm') != -1):
count_bn = count_bn + 1
for i in range(count_bn - 2):
self.insnorms.add(InstanceNorm())
with self.name_scope():
self.insnorms_fw = nn.HybridSequential(prefix='instancenormfw_')
for i in range(len(self.features._children)):
if (self.features._children[i].name.find('batchnorm') != -1):
self.insnorms_fw.add(InstanceNorm())
if self.do_countpath:
layers_drm_cp += [nn.Conv2DTranspose(channels=cfg[vgg_name][0][-2], in_channels=self.num_class, kernel_size=(1, 1), strides=(1, 1),
use_bias=False),
Reshape(shape=(self.num_class, 1, 1))]
with self.name_scope():
self.drm_cp = nn.HybridSequential(prefix='drmcp_')
for block in layers_drm_cp[::-1]:
self.drm_cp.add(block)
def __init__(self, NumLayer, filters):
super(deconv_bn_relu, self).__init__()
self.NumLayer = NumLayer
self.filters = filters
self.net = nn.HybridSequential()
with self.net.name_scope():
for i in range(NumLayer-1):
self.net.add(
nn.Conv2DTranspose(self.filters,kernel_size=4, padding=1, strides=2),
nn.InstanceNorm(),
nn.Activation('relu')
)
self.net.add(
nn.Conv2DTranspose(3, kernel_size=4, padding=1, strides=2)
)
def __init__(self, **kwargs):
super(ResnetGenerator, self).__init__(**kwargs)
num_blocks = 9
with self.name_scope():
self.module = nn.Sequential()
with self.module.name_scope():
self.module.add(nn.ReflectionPad2D(3))
self.module.add(nn.Conv2D(64, kernel_size=7, padding=0))
self.module.add(nn.InstanceNorm())
self.module.add(nn.Activation(activation='relu'))
n_downsampling = 2
for i in range(n_downsampling):
mult = 2**i
self.module.add(
nn.Conv2D(128 * mult,
kernel_size=3,
strides=2,
padding=1))
self.module.add(nn.InstanceNorm())
self.module.add(nn.Activation(activation='relu'))
for i in range(num_blocks):
self.module.add(ResnetBlock())
for i in range(n_downsampling):
mult = 2**(n_downsampling - i)
self.module.add(
nn.Conv2DTranspose(int(64 * mult / 2),
kernel_size=3,
strides=2,
padding=1,
output_padding=1))
self.module.add(nn.InstanceNorm())
self.module.add(nn.Activation(activation='relu'))
self.module.add(nn.ReflectionPad2D(3))
self.module.add(nn.Conv2D(3, kernel_size=7, padding=0))
self.module.add(nn.Activation(activation='tanh'))
def fcn_conv_2d_transpose_bn(self,
decoder,
filters,
conv_ksize,
keep_prob,
conv_strides=2,
padding=(1, 1),
activation='relu',
is_batch_norm=True,
is_training=True,
is_dropout=False):
with decoder.name_scope():
decoder.add(
nn.Conv2DTranspose(channels=filters,
kernel_size=conv_ksize,
strides=conv_strides,
padding=padding,
layout='NHWC',
activation=activation))
if is_batch_norm:
decoder.add(nn.BatchNorm(axis=3))
if is_dropout:
decoder.add(nn.Dropout(kp))
def __init__(self,
block_unit,
channels,
filter_size,
stride,
padding,
out_channels=0,
**kwargs):
super(LastTopdownBlock, self).__init__(**kwargs)
output_padding = 1 if stride > 1 else 0
layers_body = []
self.body = nn.HybridSequential(prefix='')
# THIS IS MISSING A BATCHNORM
layers_body += [
nn.BatchNorm(),
nn.Conv2DTranspose(channels=out_channels,
in_channels=channels,
kernel_size=filter_size,
strides=stride,
padding=padding,
output_padding=output_padding,
use_bias=False,
params=block_unit.body[0].params),
UpsampleLayer(size=2, scale=1.),
nn.BatchNorm()
]
for layer in layers_body[::-1]:
self.body.add(layer)
self.insnorms = nn.HybridSequential(prefix='insnorms_')
for layer in self.body:
if layer.name.find('batchnorm') != -1:
self.insnorms.add(InstanceNorm())
def _make_deconv_layer(self, num_layers, num_filters, num_kernels):
assert num_layers == len(num_filters), \
'ERROR: num_deconv_layers is different from len(num_deconv_filters)'
assert num_layers == len(num_kernels), \
'ERROR: num_deconv_layers is different from len(num_deconv_filters)'
layer = nn.HybridSequential(prefix='final_')
with layer.name_scope():
for i in range(num_layers):
kernel, padding, output_padding = \
self._get_deconv_cfg(num_kernels[i])
planes = num_filters[i]
layer.add(
nn.Conv2DTranspose(
channels=planes,
kernel_size=kernel,
strides=2,
padding=padding,
output_padding=output_padding,
use_bias=self.deconv_with_bias,
weight_initializer=initializer.Normal(0.001),
bias_initializer=initializer.Zero()))
layer.add(nn.BatchNorm(gamma_initializer=initializer.One(),
beta_initializer=initializer.Zero()))
layer.add(nn.Activation('relu'))
self.inplanes = planes
return layer
def add(self, channel):
assert(len(channel) == self.layers)
for l in range(self.layers):
self.encoder.append(nn.Conv2D(channel[l],
kernel_size = [3, 3],
strides = [2, 1],
padding = [1, 1],
dilation = [1, 1]))
conv = nn.Conv2D if 'bottleneck' in self.arch else nn.Conv2DTranspose
dchannel = 1 if l == 0 else channel[-l]
self.decoder.insert(0, conv(dchannel,
kernel_size = [3, 3],
strides = [2, 1],
padding = [1, 1],
dilation = [1, 1]))
self.register_child(self.encoder[-1])
self.register_child(self.decoder[0])
if self.reconstruct:
assert('bottleneck' not in self.arch)
assert('encoder' in self.arch and 'decoder' in self.arch)
channel = channel if l != layers - 1 else self.feature
self.reconstructor.append(nn.Conv2DTranspose(dchannel,
kernel_size = [3, 3],
strides = [2, 2],
padding = [1, 1],
dilation = [1, 1]))
self.register_child(self.reconstructor[-1])
if self.norm:
self.enorm.append(self.block['norm'](axis = 2))
self.dnorm.append(self.block['norm'](axis = 2))
self.register_child(self.enorm[-1])
self.register_child(self.dnorm[-1])
if self.reconstruct:
self.rnorm.append(self.block['norm'](axis = 2))
self.register_child(self.rnorm[-1])
def __init__(self, dim2, classes, out_size, bn_epsilon, **kwargs):
super(SBDecoder, self).__init__(**kwargs)
with self.name_scope():
self.decode1 = SBDecodeBlock(
channels=classes,
out_size=((out_size[0] // 8,
out_size[1] // 8) if out_size else None),
bn_epsilon=bn_epsilon)
self.decode2 = SBDecodeBlock(
channels=classes,
out_size=((out_size[0] // 4,
out_size[1] // 4) if out_size else None),
bn_epsilon=bn_epsilon)
self.conv3c = conv1x1_block(in_channels=dim2,
out_channels=classes,
bn_epsilon=bn_epsilon,
activation=(lambda: PReLU2(classes)))
self.output = nn.Conv2DTranspose(channels=classes,
kernel_size=2,
strides=2,
padding=0,
output_padding=0,
in_channels=classes,
use_bias=False)
self.up = InterpolationBlock(scale_factor=2, out_size=out_size)
def test_gradient():
ctx = mx.cpu()
num_latent_maps = 1
input_shape = (1, 2, 2)
input_dim = 4
batch_size = 2
# build the network
enc_nn = nn.HybridSequential()
enc_nn.add(nn.Conv2D(channels=2, kernel_size=(1, 1), activation='relu', bias_initializer=mx.init.Uniform(1.0)))
dec_nn = nn.HybridSequential()
dec_nn.add(nn.Conv2DTranspose(channels=1, kernel_size=(1, 1), bias_initializer=mx.init.Uniform(1.0)))
conv_draw_nn = ConvDRAW(enc_nn, dec_nn, num_steps=2, batch_size=batch_size, input_shape=input_shape,
num_latent_maps=num_latent_maps, encoder_output_shape=(2, 2, 2), rnn_hidden_channels=1,
kernel_size=(1, 1), ctx=ctx)
model_params = conv_draw_nn.collect_params()
mx.random.seed(np.random.randint(1000000))
model_params.initialize(init=mx.init.Uniform(1.0), ctx=ctx) # don't initialize to small weights
# loss function
loss_fn = ConvDRAWLoss(SigmoidBinaryCrossEntropyLoss(from_sigmoid=False, batch_axis=0), input_dim, (1, 2, 2))
def fwd(x):
y, q, p = conv_draw_nn(x)
return nd.sum(loss_fn(x, q, p, y))
batch_x = mx.nd.random_uniform(shape=(batch_size, *input_shape))
fwd(batch_x) # the following check fails for the first parameter if fwd is not called at least once before it.
for p in model_params.values():
assert check_gradient(fwd, [batch_x], p)
def __init__(self, pretrained=False):
super(DeepLab, self).__init__()
with self.name_scope():
resnet = gluon.model_zoo.vision.resnet101_v1(
pretrained=pretrained).features
self.layer0 = nn.HybridSequential()
self.layer0.add(*resnet[0:4])
self.layer1 = resnet[4]
self.layer2 = resnet[5]
self.layer3 = resnet[6]
self.layer4 = resnet[7]
for i in range(len(self.layer4)):
self.layer4[i].body[3]._kwargs['stride'] = (1, 1)
self.layer4[i].body[3]._kwargs['dilate'] = (2, 2)
self.layer4[i].body[3]._kwargs['pad'] = (2, 2)
self.layer4[0].downsample[0]._kwargs['stride'] = (1, 1)
self.aspp = ASPP(512)
self.pred = nn.HybridSequential()
self.pred.add(nn.Conv2D(512, kernel_size=1), nn.BatchNorm(),
nn.Activation('relu'),
nn.Conv2D(cfg.n_class, kernel_size=1))
self.refine = BoundaryRefineModule(cfg.n_class)
self.upsample = nn.Conv2DTranspose(cfg.n_class,
kernel_size=32,
strides=16,
padding=8)
def _make_deconv_layer(self, num_filters, num_kernels):
"""Make deconv layers using the configs"""
assert len(num_kernels) == len(num_filters), \
'Deconv filters and kernels number mismatch: {} vs. {}'.format(
len(num_filters), len(num_kernels))
layers = nn.HybridSequential('deconv_')
with warnings.catch_warnings(record=True) as w:
warnings.simplefilter("always")
self.base_network.initialize()
in_planes = self.base_network(mx.nd.zeros((1, 3, 256, 256))).shape[1]
for planes, k in zip(num_filters, num_kernels):
kernel, padding, output_padding = self._get_deconv_cfg(k)
layers.add(nn.Conv2D(channels=planes,
kernel_size=3,
strides=1,
padding=1,
in_channels=in_planes))
layers.add(nn.BatchNorm())
layers.add(nn.Activation('relu'))
layers.add(nn.Conv2DTranspose(channels=planes,
kernel_size=kernel,
strides=2,
padding=padding,
output_padding=output_padding,
use_bias=False,
in_channels=planes,
weight_initializer=BilinearUpSampleInitializer()))
layers.add(nn.BatchNorm())
layers.add(nn.Activation('relu'))
in_planes = planes
return layers
def _make_deconv_layers(self):
num_layers = 3
num_filters = [256, 128, 64]
num_kernels = [4, 4, 4]
layers = nn.HybridSequential()
for i in range(num_layers):
num_filter = num_filters[i]
kernel_size = num_kernels[i]
# kernel, padding, output_padding = \
# self._get_deconv_cfg(num_kernels[i], i)
layers.add(
nn.Conv2D(channels=num_filter,
kernel_size=3,
strides=1,
padding=1,
dilation=1,
use_bias=False))
layers.add(nn.BatchNorm())
layers.add(nn.Activation(activation='relu'))
layers.add(
nn.Conv2DTranspose(channels=num_filter,
kernel_size=kernel_size,
strides=(2, 2),
padding=(1, 1),
output_padding=(0, 0)))
layers.add(nn.BatchNorm())
layers.add(nn.Activation(activation='relu'))
return layers
