def main(args):
# Load the data (DataLoader object)
path_monnet = args.Monet_Path
path_pictures = args.Pictures_Path
save_path = args.Save_Path
batch_size = args.batch_size
n_epochs = args.epochs
device = args.device
dataset = get_data_loader(path_monnet, path_pictures, batch_size)
# Create Generators and Discriminators and put them on GPU/TPU
generator_AB = Generator().to(device)
generator_BA = Generator().to(device)
discriminator_A = Discriminator().to(device)
discriminator_B = Discriminator().to(device)
generator_AB.apply(weights_init_normal)
generator_BA.apply(weights_init_normal)
discriminator_A.apply(weights_init_normal)
discriminator_B.apply(weights_init_normal)
# Set optimizers
G_optimizer = torch.optim.Adam(itertools.chain(generator_AB.parameters(),
generator_BA.parameters()),
lr=2e-4)
D_optimizer = torch.optim.Adam(itertools.chain(
discriminator_A.parameters(), discriminator_B.parameters()),
lr=2e-4)
# Set trainer
trainer = Trainer(
generator_ab=generator_AB,
generator_ba=generator_BA,
discriminator_a=discriminator_A,
discriminator_b=discriminator_B,
generator_optimizer=G_optimizer,
discriminator_optimizer=D_optimizer,
n_epochs=n_epochs,
dataloader=dataset,
device=device,
)
# Launch Training
trainer.train()
# Save the model and the loss during training
# Save logs
trainer.log.save(os.path.join(save_path, 'save_loss.txt'))
# Save the model
torch.save(generator_AB.state_dict(),
os.path.join(save_path, 'generator_AB.pt'))
torch.save(generator_BA.state_dict(),
os.path.join(save_path, 'generator_BA.pt'))
torch.save(discriminator_A.state_dict(),
os.path.join(save_path, 'discriminator_A.pt'))
torch.save(discriminator_B.state_dict(),
os.path.join(save_path, 'discriminator_B.pt'))
classname = m.__class__.__name__
if classname.find('Conv') != -1:
m.weight.data.normal_(0.0, 0.02)
elif classname.find('BatchNorm') != -1:
m.weight.data.normal_(1.0, 0.02)
m.bias.data.fill_(0)
netG = Generator(ngpu).to(device)
netG.apply(weights_init)
# load weights to test the model
# netG.load_state_dict(torch.load('weights/netG_epoch_24.pth'))
print(netG)
netD = Discriminator(ngpu).to(device)
netD.apply(weights_init)
# load weights to test the model
# netD.load_state_dict(torch.load('weights/netD_epoch_24.pth'))
print(netD)
criterion = nn.BCELoss()
# setup optimizer
optimizerD = optim.Adam(netD.parameters(), lr=0.0001, betas=(0.5, 0.999))
optimizerG = optim.Adam(netG.parameters(), lr=0.0001, betas=(0.5, 0.999))
fixed_noise = torch.randn(128, nz, 1, 1, device=device)
real_label = 1
fake_label = 0
niter = 400