def set_network(depth, ctx, lr, beta1, ndf, ngf, append=True, solver='adam'):
# Pixel2pixel networks
if append:
netD = models.Discriminator(
in_channels=6, n_layers=2, ndf=ndf
) ##netG = models.CEGenerator(in_channels=3, n_layers=depth, ndf=ngf) # UnetGenerator(in_channels=3, num_downs=8) #
else:
netD = models.Discriminator(in_channels=3, n_layers=2, ndf=ndf)
#netG = models.UnetGenerator(in_channels=3, num_downs =depth, ngf=ngf) # UnetGenerator(in_channels=3, num_downs=8) #
#netD = models.Discriminator(in_channels=6, n_layers =depth-1, ndf=ngf/4)
netG = models.CEGenerator(in_channels=3, n_layers=depth, ndf=ngf)
# Initialize parameters
models.network_init(netG, ctx=ctx)
models.network_init(netD, ctx=ctx)
if solver == 'adam':
# trainer for the generator and the discriminator
trainerG = gluon.Trainer(netG.collect_params(), 'adam', {
'learning_rate': lr,
'beta1': beta1
})
trainerD = gluon.Trainer(netD.collect_params(), 'adam', {
'learning_rate': lr,
'beta1': beta1
})
elif solver == 'sgd':
print('sgd')
trainerG = gluon.Trainer(netG.collect_params(), 'sgd', {
'learning_rate': lr,
'momentum': 0.9
})
trainerD = gluon.Trainer(netD.collect_params(), 'sgd', {
'learning_rate': lr,
'momentum': 0.9
})
return netG, netD, trainerG, trainerD
def set_network(depth, ctx, lr, beta1, ndf, ngf, append=True):
if append:
netD = models.Discriminator(in_channels=6,
n_layers=depth - 1,
istest=True,
ndf=ndf)
else:
netD = models.Discriminator(in_channels=3,
n_layers=depth - 1,
istest=True,
ndf=ndf)
netG = models.CEGenerator(
in_channels=3, n_layers=depth, istest=True,
ndf=ngf) # UnetGenerator(in_channels=3, num_downs=8) #
# Initialize parameters
models.network_init(netG, ctx=ctx)
models.network_init(netD, ctx=ctx)
# trainer for the generator and the discriminator
trainerG = gluon.Trainer(netG.collect_params(), 'adam', {
'learning_rate': lr,
'beta1': beta1
})
trainerD = gluon.Trainer(netD.collect_params(), 'adam', {
'learning_rate': lr,
'beta1': beta1
})
return netG, netD, trainerG, trainerD
def set_network():
# Pixel2pixel networks
netG = models.CEGenerator(
in_channels=3) # UnetGenerator(in_channels=3, num_downs=8) #
netD = models.Discriminator(in_channels=6)
# Initialize parameters
models.network_init(netG, ctx=ctx)
models.network_init(netD, ctx=ctx)
# trainer for the generator and the discriminator
trainerG = gluon.Trainer(netG.collect_params(), 'adam', {
'learning_rate': lr,
'beta1': beta1
})
trainerD = gluon.Trainer(netD.collect_params(), 'adam', {
'learning_rate': lr,
'beta1': beta1
})
return netG, netD, trainerG, trainerD
