def main(args): train_loader, test_loader = load_data(args) GeneratorA2B = CycleGAN() GeneratorB2A = CycleGAN() DiscriminatorA = Discriminator() DiscriminatorB = Discriminator() if args.cuda: GeneratorA2B = GeneratorA2B.cuda() GeneratorB2A = GeneratorB2A.cuda() DiscriminatorA = DiscriminatorA.cuda() DiscriminatorB = DiscriminatorB.cuda() optimizerG = optim.Adam(itertools.chain(GeneratorA2B.parameters(), GeneratorB2A.parameters()), lr=args.lr, betas=(0.5, 0.999)) optimizerD = optim.Adam(itertools.chain(DiscriminatorA.parameters(), DiscriminatorB.parameters()), lr=args.lr, betas=(0.5, 0.999)) if args.training: path = 'E:/cyclegan/checkpoints/model_{}_{}.pth'.format(285, 200) checkpoint = torch.load(path) GeneratorA2B.load_state_dict(checkpoint['generatorA']) GeneratorB2A.load_state_dict(checkpoint['generatorB']) DiscriminatorA.load_state_dict(checkpoint['discriminatorA']) DiscriminatorB.load_state_dict(checkpoint['discriminatorB']) optimizerG.load_state_dict(checkpoint['optimizerG']) optimizerD.load_state_dict(checkpoint['optimizerD']) start_epoch = 285 else: init_net(GeneratorA2B, init_type='normal', init_gain=0.02, gpu_ids=[0]) init_net(GeneratorB2A, init_type='normal', init_gain=0.02, gpu_ids=[0]) init_net(DiscriminatorA, init_type='normal', init_gain=0.02, gpu_ids=[0]) init_net(DiscriminatorB, init_type='normal', init_gain=0.02, gpu_ids=[0]) start_epoch = 1 if args.evaluation: evaluation(test_loader, GeneratorA2B, GeneratorB2A, args) else: cycle = nn.L1Loss() gan = nn.BCEWithLogitsLoss() identity = nn.L1Loss() for epoch in range(start_epoch, args.epochs): train(train_loader, GeneratorA2B, GeneratorB2A, DiscriminatorA, DiscriminatorB, optimizerG, optimizerD, cycle, gan, identity, args, epoch) evaluation(test_loader, GeneratorA2B, GeneratorB2A, args)
# configure full paths checkpoints_dir = os.path.join(checkpoints_dir, experiment_id) samples_dir = os.path.join(samples_dir, experiment_id) logs_dir = os.path.join(logs_dir, experiment_id) # make directories os.system('mkdir -p ' + checkpoints_dir) os.system('mkdir -p ' + samples_dir) os.system('mkdir -p ' + logs_dir) ## create models # call the function to get models G_XtoY, G_YtoX, Dp_X, Dp_Y, Dg_X, Dg_Y = CycleGAN(n_res_blocks=2) # define optimizer parameters g_params = list(G_XtoY.parameters()) + list( G_YtoX.parameters()) # Get generator parameters # Create optimizers for the generators and discriminators g_optimizer = optim.Adam(g_params, g_lr, [beta1, beta2]) dp_x_optimizer = optim.Adam(Dp_X.parameters(), d_lr, [beta1, beta2]) dp_y_optimizer = optim.Adam(Dp_Y.parameters(), d_lr, [beta1, beta2]) dg_x_optimizer = optim.Adam(Dg_X.parameters(), d_lr, [beta1, beta2]) dg_y_optimizer = optim.Adam(Dg_Y.parameters(), d_lr, [beta1, beta2]) # count number of parameters in a model def count_model_parameters(model): n_params = sum(p.numel() for p in model.parameters()) n_trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)