def model_create(self):
from .. import models_64x64
config = self.myargs.config.trainer.model
if config.d_use_lp:
self.logger.info("Use LayerNorm in D.")
D = models_64x64.DiscriminatorWGANGPLN(3)
else:
self.logger.info("Use InstanceNorm in D.")
D = models_64x64.DiscriminatorWGANGP(3)
G = models_64x64.Generator(config.z_dim)
G_ema = models_64x64.Generator(config.z_dim)
ema = ema_model.EMA(G, G_ema, decay=0.9999, start_itr=config.ema_start)
D.cuda()
G.cuda()
G_ema.cuda()
self.myargs.checkpoint_dict['D'] = D
self.myargs.checkpoint_dict['G'] = G
self.myargs.checkpoint_dict['G_ema'] = G_ema
self.logger.info_msg('Number of params in G: {} D: {}'.format(*[
sum([p.data.nelement() for p in net.parameters()])
for net in [G, D]
]))
return G, D, G_ema, ema
def model_create(self):
from .. import models_64x64
config = self.myargs.config.model
D = models_64x64.DiscriminatorWGANGP(3)
print(D)
G = models_64x64.Generator(config.z_dim)
print(G)
G_ema = models_64x64.Generator(config.z_dim)
ema = ema_model.EMA(G, G_ema, decay=0.9999, start_itr=config.ema_start)
D = D.cuda(self.args.rank)
G = G.cuda(self.args.rank)
G_ema = G_ema.cuda(self.args.rank)
self.myargs.checkpoint_dict['D'] = D
self.myargs.checkpoint_dict['G'] = G
self.myargs.checkpoint_dict['G_ema'] = G_ema
print('Number of params in G: {} D: {}'.format(*[
sum([p.data.nelement() for p in net.parameters()])
for net in [G, D]
]))
return G, D, G_ema, ema
def model_create(self):
from .. import models_64x64
config = self.myargs.config.model
D = models_64x64.DiscriminatorWGANGP(3)
print(D)
G = models_64x64.Generator(config.z_dim)
print(G)
G_ema = models_64x64.Generator(config.z_dim)
ema = ema_model.EMA(G, G_ema, decay=0.9999, start_itr=config.ema_start)
D.cuda()
G.cuda()
G_ema.cuda()
if config.parallel:
D = nn.DataParallel(D)
G = nn.DataParallel(G)
self.myargs.checkpoint_dict['D'] = D
self.myargs.checkpoint_dict['G'] = G
self.myargs.checkpoint_dict['G_ema'] = G_ema
self.print_number_params(models=dict(G=G, D=D))
return G, D, G_ema, ema
def model_create(self):
config = self.config.model
resolution = imsize_dict[self.config.dataset.dataset]
n_classes = nclass_dict[self.config.dataset.dataset]
self.n_classes = n_classes
if hasattr(self.config.dataset, 'attr'):
assert n_classes == 2**len(self.config.dataset.attr)
if getattr(config, 'use_cbn', False):
self.G = SharedGeneratorCBN(resolution=resolution,
n_classes=n_classes,
no_optim=False,
config=config.generator).cuda()
else:
self.G = SharedGeneratorNoSkip(resolution=resolution,
no_optim=False,
config=config.generator).cuda()
self.G_optim = self.G.optim
self.myargs.checkpoint_dict['G'] = self.G
self.myargs.checkpoint_dict['G_optim'] = self.G.optim
if getattr(config, 'use_ema', False):
self.G_ema = copy.deepcopy(self.G)
self.myargs.checkpoint_dict['G_ema'] = self.G_ema
self.ema = ema_model.EMA(self.G,
self.G_ema,
decay=config.ema_decay,
start_itr=config.ema_start)
# Create controller
controller_c = config.controller
self.controller = Controller(
n_classes=n_classes,
num_layers=(self.G.num_layers +
1 if self.G.output_sample_arc else self.G.num_layers),
num_branches=len(self.G.ops),
config=controller_c).cuda()
self.controller_optim = self.controller.optim
self.myargs.checkpoint_dict['controller'] = self.controller
self.myargs.checkpoint_dict['C_optim'] = self.controller.optim
self.G_C = G_Controller(self.G, self.controller)
# self.G_C = torch.nn.DataParallel(self.G_C)
if getattr(config, 'use_cdisc', False):
disc_c = config.discriminator_cond
D_activation = activation_dict[disc_c.D_activation]
self.D = BigGAN.Discriminator(
**{
**disc_c, 'resolution': resolution,
'n_classes': n_classes,
'D_activation': D_activation
}, **disc_c.optimizer).cuda()
self.D_optim = self.D.optim
else:
from ..models.autogan_cifar10_a import Discriminator
from ..models import optimizer_dict
disc_c = config.discriminator
disc_optim_c = disc_c.optimizer
self.D = Discriminator(args=disc_c).cuda()
adam_eps = getattr(disc_optim_c, 'adam_eps', 1.e-8)
self.D_optim = self.optim = optimizer_dict[disc_optim_c.type](
params=self.D.parameters(),
lr=disc_optim_c.D_lr,
betas=(disc_optim_c.D_B1, disc_optim_c.D_B2),
eps=adam_eps)
self.myargs.checkpoint_dict['D'] = self.D
self.myargs.checkpoint_dict['D_optim'] = self.D_optim
self.D = torch.nn.DataParallel(self.D)
self.models = {'controller': self.controller, 'G': self.G, 'D': self.D}
self.print_number_params(models=self.models)
def model_create(self):
import BigGAN as model
import utils
config = self.config.model
Generator = model.Generator
Discriminator = model.Discriminator
G_D = model.G_D
print('Create generator: {}'.format(Generator))
self.resolution = utils.imsize_dict[self.config.dataset.dataset]
self.n_classes = utils.nclass_dict[self.config.dataset.dataset]
G_activation = utils.activation_dict[config.Generator.G_activation]
self.G = Generator(
**{
**config.Generator, 'resolution': self.resolution,
'n_classes': self.n_classes,
'G_activation': G_activation
}, **config.optimizer).cuda()
optim_type = getattr(config.optimizer, 'optim_type', None)
if optim_type == 'radam':
print('Using radam optimizer.')
from template_lib.optimizer import radam
self.G.optim = radam.RAdam(params=self.G.parameters(),
lr=config.optimizer.G_lr,
betas=(config.optimizer.G_B1,
config.optimizer.G_B2),
weight_decay=0,
eps=config.optimizer.adam_eps)
print('Create discriminator: {}'.format(Discriminator))
D_activation = utils.activation_dict[config.Discriminator.D_activation]
self.D = Discriminator(logger=self.logger,
**{
**config.Discriminator, 'resolution':
self.resolution,
'n_classes': self.n_classes,
'D_activation': D_activation
},
**config.optimizer).cuda()
if optim_type == 'radam':
print('Using radam optimizer.')
from TOOLS import radam
self.D.optim = radam.RAdam(params=self.D.parameters(),
lr=config.optimizer.D_lr,
betas=(config.optimizer.D_B1,
config.optimizer.D_B2),
weight_decay=0,
eps=config.optimizer.adam_eps)
if getattr(self.config.train_one_epoch, 'weigh_loss', False):
self.create_alpha()
self.G_ema = Generator(logger=self.logger,
**{
**config.Generator, 'resolution':
self.resolution,
'n_classes': self.n_classes,
'G_activation': G_activation,
'skip_init': True,
'no_optim': True
}).cuda()
self.ema = ema_model.EMA(self.G,
self.G_ema,
decay=0.9999,
start_itr=config.ema_start)
print('Create G_D: {}'.format(G_D))
self.GD = G_D(self.G, self.D)
if config['parallel']:
self.GD = nn.DataParallel(self.GD)
self.myargs.checkpoint_dict['G'] = self.G
self.myargs.checkpoint_dict['G_optim'] = self.G.optim
self.myargs.checkpoint_dict['D'] = self.D
self.myargs.checkpoint_dict['D_optim'] = self.D.optim
self.myargs.checkpoint_dict['G_ema'] = self.G_ema
models = {'G': self.G, 'D': self.D}
self.print_number_params(models=models)