def build_optimizers(generator, discriminator, config):
optimizer = config['training']['optimizer']
lr_g = config['training']['lr_g']
lr_d = config['training']['lr_d']
equalize_lr = config['training']['equalize_lr']
toggle_grad(generator, True)
toggle_grad(discriminator, True)
if equalize_lr:
g_gradient_scales = getattr(generator, 'gradient_scales', dict())
d_gradient_scales = getattr(discriminator, 'gradient_scales', dict())
g_params = get_parameter_groups(generator.parameters(),
g_gradient_scales,
base_lr=lr_g)
d_params = get_parameter_groups(discriminator.parameters(),
d_gradient_scales,
base_lr=lr_d)
else:
g_params = generator.parameters()
d_params = discriminator.parameters()
# Optimizers
if optimizer == 'rmsprop':
g_optimizer = optim.RMSprop(g_params, lr=lr_g, alpha=0.99, eps=1e-8)
d_optimizer = optim.RMSprop(d_params, lr=lr_d, alpha=0.99, eps=1e-8)
elif optimizer == 'adam':
g_optimizer = optim.Adam(g_params, lr=lr_g, betas=(0., 0.99), eps=1e-8)
d_optimizer = optim.Adam(d_params, lr=lr_d, betas=(0., 0.99), eps=1e-8)
elif optimizer == 'sgd':
g_optimizer = optim.SGD(g_params, lr=lr_g, momentum=0.)
d_optimizer = optim.SGD(d_params, lr=lr_d, momentum=0.)
return g_optimizer, d_optimizer
def build_optimizers(generator, discriminator, config):
optimizer = config['training']['optimizer']
lr_g = config['training']['lr_g']
lr_d = config['training']['lr_d']
toggle_grad(generator, True)
toggle_grad(discriminator, True)
g_params = generator.parameters()
d_params = discriminator.parameters()
if optimizer == 'rmsprop':
g_optimizer = optim.RMSprop(g_params, lr=lr_g, alpha=0.99, eps=1e-8)
d_optimizer = optim.RMSprop(d_params, lr=lr_d, alpha=0.99, eps=1e-8)
elif optimizer == 'adam':
beta1 = config['training']['beta1']
beta2 = config['training']['beta2']
g_optimizer = optim.Adam(g_params,
lr=lr_g,
betas=(beta1, beta2),
eps=1e-8)
d_optimizer = optim.Adam(d_params,
lr=lr_d,
betas=(beta1, beta2),
eps=1e-8)
elif optimizer == 'sgd':
g_optimizer = optim.SGD(g_params, lr=lr_g, momentum=0.)
d_optimizer = optim.SGD(d_params, lr=lr_d, momentum=0.)
return g_optimizer, d_optimizer