def run(args): # Get lasagne weights lasagne_weights_path = args.weights print('Loading lasagne weights') with open(lasagne_weights_path, "rb") as f: G, D, Gs = cPickle.load(f) # for l in lasagne.layers.get_all_layers(D.find_layer('Dscoresws')): # print('%-12s %-30s %d' % (l.name, `lasagne.layers.get_output_shape(l)`, lasagne.layers.count_params(l))) if args.generator: print('Converting Generator model') generator = Generator() convert_generator(Gs, generator) _, model_name = os.path.split(args.weights) model_name = model_name.replace('.pkl', '.pth') output_path = os.path.join(args.generator, model_name) print('Saving model to {}'.format(output_path)) torch.save(generator.state_dict(), output_path) if args.discriminator: print('Converting Discriminator model') discriminator = Discriminator() convert_discriminator(D, discriminator) _, model_name = os.path.split(args.weights) model_name = model_name.replace('.pkl', '.pth') output_path = os.path.join(args.discriminator, model_name) print('Saving model to {}'.format(output_path)) torch.save(discriminator.state_dict(), output_path)
class gan(nn.Module): # def __init__(self, params, save_dir, g_weight_dir, d_weight_dir, d_update_freq=1, start_epoch=0, g_lr=2e-4, d_lr=2e-4, use_cuda=True): def __init__(self, params, args): super(gan, self).__init__() self.G = MModel(params, use_cuda=True) self.D = Discriminator(params, bias=True) self.vgg_loss = VGGPerceptualLoss() self.L1_loss = nn.L1Loss() if args.use_cuda: self.G = self.G.cuda() self.D = self.D.cuda() self.vgg_loss = self.vgg_loss.cuda() self.L1_loss = self.L1_loss.cuda() if args.g_weight_dir: self.G.load_state_dict(torch.load(args.g_weight_dir), strict=True) if args.d_weight_dir: self.D.load_state_dict(torch.load(args.d_weight_dir), strict=False) self.optimizer_G = torch.optim.Adam(self.G.parameters(), lr=args.g_lr) self.optimizer_D = torch.optim.Adam(self.D.parameters(), lr=args.d_lr) self.save_dir = args.save_dir if not os.path.exists(args.save_dir): os.makedirs(args.save_dir) self.d_update_freq = args.d_update_freq self.save_freq = args.save_freq self.writer = SummaryWriter('runs/' + args.save_dir) self.use_cuda = args.use_cuda self.start_epoch = args.start_epoch def G_loss(self, input, target): vgg = self.vgg_loss(input, target) L1 = self.L1_loss(input, target) # return vgg return vgg + L1 def update_D(self, loss, epoch): if epoch % self.d_update_freq == 0: loss.backward() self.optimizer_D.step() def get_patch_weight(self, pose, size=62): heads = pose[:, 0, :, :] heads = heads.unsqueeze(1) heads = torch.nn.functional.interpolate(heads, size=size) heads = heads * 5 + torch.ones_like(heads) return heads def gan_loss(self, out, label, pose): # weight = self.get_patch_weight(pose) # return nn.BCELoss(weight=weight)(out, torch.ones_like(out) if label==1 else torch.zeros_like(out)) return nn.BCELoss()( out, torch.ones_like(out) if label == 1 else torch.zeros_like(out)) def train(self, dl, epoch): # i -- current epoch cnt = 0 loss_D_real_sum, loss_D_fake_sum, loss_D_sum, loss_G_gan_sum, loss_G_img_sum, loss_G_sum = 0, 0, 0, 0, 0, 0 for iter, (src_img, y, src_pose, tgt_pose, src_mask_prior, x_trans, src_mask_gt, tgt_face, tgt_face_box, src_face_box) in enumerate(dl): print('epoch:', epoch, 'iter:', iter) self.optimizer_D.zero_grad() if self.use_cuda: src_img, y, src_pose, tgt_pose, src_mask_prior, x_trans = src_img.cuda( ), y.cuda(), src_pose.cuda(), tgt_pose.cuda( ), src_mask_prior.cuda(), x_trans.cuda() out = self.G(src_img, src_pose, tgt_pose, src_mask_prior, x_trans) gen = out[0] loss_D_real = self.gan_loss(self.D(y, tgt_pose), 1, tgt_pose) loss_D_fake = self.gan_loss(self.D(gen.detach(), tgt_pose), 0, tgt_pose) loss_D = loss_D_real + loss_D_fake self.update_D(loss_D, epoch) if False and epoch < 10: loss_G_gan = torch.zeros((1)) loss_G_img = torch.zeros((1)) loss_G = loss_G_gan + loss_G_img else: self.optimizer_G.zero_grad() loss_G_gan = self.gan_loss(self.D(gen, tgt_pose), 1, tgt_pose) loss_G_img = self.G_loss(gen, y) # vgg_loss + L1_loss loss_G = loss_G_gan + loss_G_img loss_G.backward() self.optimizer_G.step() loss_D_real_sum += loss_D_real.item() loss_D_fake_sum += loss_D_fake.item() loss_D_sum += loss_D.item() loss_G_gan_sum += loss_G_gan.item() loss_G_img_sum += loss_G_img.item() loss_G_sum += loss_G.item() cnt += 1 # if epoch % self.save_freq == 0 and iter < 3: # self.writer.add_images('gen/epoch%d'%epoch, gen*0.5+0.5) # self.writer.add_images('y/epoch%d'%epoch, y*0.5+0.5) # self.writer.add_images('src_mask/epoch%d'%epoch, out[2].view((out[2].size(0)*out[2].size(1), 1, out[2].size(2), out[2].size(3)))) # self.writer.add_images('warped/epoch%d'%epoch, out[3].view((out[3].size(0)*11, 3, out[3].size(2), out[3].size(3)))*0.5+0.5) self.writer.add_scalar('loss_D_real', loss_D_real_sum / cnt, epoch) self.writer.add_scalar('loss_D_fake', loss_D_fake_sum / cnt, epoch) self.writer.add_scalar('loss_D', loss_D_sum / cnt, epoch) self.writer.add_scalar('loss_G_gan', loss_G_gan_sum / cnt, epoch) self.writer.add_scalar('loss_G_img', loss_G_img_sum / cnt, epoch) self.writer.add_scalar('loss_G', loss_G_sum / cnt, epoch) self.writer.add_scalars('DG', { 'D': loss_D / cnt, 'G': loss_G / cnt }, epoch) if epoch % self.save_freq == 0: torch.save(self.G.state_dict(), os.path.join(self.save_dir, 'g_epoch_%d.pth' % epoch)) torch.save(self.D.state_dict(), os.path.join(self.save_dir, 'd_epoch_%d.pth' % epoch)) def test(self, test_dl, epoch): self.G.eval() for iter, (src_img, y, src_pose, tgt_pose, src_mask_prior, x_trans, src_mask_gt, tgt_face, tgt_face_box, src_face_box) in enumerate(test_dl): print('test', 'epoch:', epoch, 'iter:', iter) if self.use_cuda: src_img, y, src_pose, tgt_pose, src_mask_prior, x_trans = src_img.cuda( ), y.cuda(), src_pose.cuda(), tgt_pose.cuda( ), src_mask_prior.cuda(), x_trans.cuda() with torch.no_grad(): out = self.G(src_img, src_pose, tgt_pose, src_mask_prior, x_trans) gen = out[0] if iter == 0: self.writer.add_images('test_gen/epoch%d' % epoch, gen * 0.5 + 0.5) self.writer.add_images('test_y/epoch%d' % epoch, y * 0.5 + 0.5) self.writer.add_images('test_src/epoch%d' % epoch, src_img * 0.5 + 0.5) self.writer.add_images( 'test_src_mask/epoch%d' % epoch, out[2].view( (out[2].size(0) * out[2].size(1), 1, out[2].size(2), out[2].size(3))))
class AdvGAN_Pretrain: def __init__(self, device, model, model_num_labels, box_min, box_max): self.device = device self.model_num_labels = model_num_labels self.model = model self.box_min = box_min self.box_max = box_max self.netG = Generator().to(device) self.netDisc = Discriminator().to(device) # initialize all weights self.netG.apply(weights_init) self.netDisc.apply(weights_init) # initialize optimizers self.optimizer_G = torch.optim.Adam(self.netG.parameters(), lr=1e-3) self.optimizer_D = torch.optim.Adam(self.netDisc.parameters(), lr=1e-3) if not os.path.exists(models_path): os.makedirs(models_path) def train_batch(self, x, labels): # optimize D for i in range(1): perturbation = self.netG(x) # add a clipping trick adv_images = torch.clamp(perturbation, -0.3, 0.3) + x adv_images = torch.clamp(adv_images, self.box_min, self.box_max) self.optimizer_D.zero_grad() pred_real = self.netDisc(x) loss_D_real = F.mse_loss(pred_real, torch.ones_like(pred_real, device=self.device)) loss_D_real.backward() pred_fake = self.netDisc(adv_images.detach()) loss_D_fake = F.mse_loss(pred_fake, torch.zeros_like(pred_fake, device=self.device)) loss_D_fake.backward() loss_D_GAN = loss_D_fake + loss_D_real self.optimizer_D.step() # optimize G for i in range(1): self.optimizer_G.zero_grad() # cal G's loss in GAN pred_fake = self.netDisc(adv_images) loss_G_fake = F.mse_loss(pred_fake, torch.ones_like(pred_fake, device=self.device)) loss_G_fake.backward(retain_graph=True) # calculate perturbation norm loss_perturb = torch.mean(torch.norm(perturbation.view(perturbation.shape[0], -1), 2, dim=1)) # loss_perturb = torch.max(loss_perturb - C, torch.zeros(1, device=self.device)) # cal adv loss logits_model = self.model(adv_images) probs_model = F.softmax(logits_model, dim=1) onehot_labels = torch.eye(self.model_num_labels, device=self.device)[labels] # C&W loss function real = torch.sum(onehot_labels * probs_model, dim=1) other, _ = torch.max((1 - onehot_labels) * probs_model - onehot_labels * 10000, dim=1) zeros = torch.zeros_like(other) loss_adv_arr = torch.max(real - other, zeros) print(loss_adv_arr) print(loss_adv_arr.shape) loss_adv = torch.sum(loss_adv) # maximize cross_entropy loss # loss_adv = -F.mse_loss(logits_model, onehot_labels) # loss_adv = - F.cross_entropy(logits_model, labels) adv_lambda = 10 pert_lambda = 1 loss_G = adv_lambda * loss_adv + pert_lambda * loss_perturb loss_G.backward() self.optimizer_G.step() return loss_D_GAN.item(), loss_G_fake.item(), loss_perturb.item(), loss_adv.item() def train(self, train_dataloader, epochs): writer = SummaryWriter(log_dir="visualization/orig_advgan/", comment='Original AdvGAN stats') for epoch in range(1, epochs+1): if epoch == 50: self.optimizer_G = torch.optim.Adam(self.netG.parameters(), lr=1e-4) self.optimizer_D = torch.optim.Adam(self.netDisc.parameters(), lr=1e-4) if epoch == 80: self.optimizer_G = torch.optim.Adam(self.netG.parameters(), lr=1e-5) self.optimizer_D = torch.optim.Adam(self.netDisc.parameters(), lr=1e-5) loss_D_sum = 0 loss_G_fake_sum = 0 loss_perturb_sum = 0 loss_adv_sum = 0 for i, data in enumerate(train_dataloader, start=0): images, labels = data images, labels = images.to(self.device), labels.to(self.device) loss_D_batch, loss_G_fake_batch, loss_perturb_batch, loss_adv_batch = \ self.train_batch(images, labels) loss_D_sum += loss_D_batch loss_G_fake_sum += loss_G_fake_batch loss_perturb_sum += loss_perturb_batch loss_adv_sum += loss_adv_batch # print statistics num_batch = len(train_dataloader) writer.add_scalar('discriminator_loss', loss_D_sum/num_batch, epoch) writer.add_scalar('generator_loss', loss_G_fake_sum/num_batch, epoch) writer.add_scalar('perturbation_loss', loss_perturb_sum/num_batch, epoch) writer.add_scalar('adversarial_loss', loss_adv_sum/num_batch, epoch) print("epoch %d:\nloss_D: %.5f, loss_G_fake: %.5f,\ \nloss_perturb: %.5f, loss_adv: %.5f\n" % (epoch, loss_D_sum/num_batch, loss_G_fake_sum/num_batch, loss_perturb_sum/num_batch, loss_adv_sum/num_batch)) # save generator if epoch%20==0: netG_file_name = models_path + 'netG_original_epoch_' + str(epoch) + '.pth' torch.save(self.netG.state_dict(), netG_file_name) netDisc_file_name = models_path + 'netDisc_original_epoch_' + str(epoch) + '.pth' torch.save(self.netDisc.state_dict(), netDisc_file_name) writer.close()
accumulate, generator.generate_7, inference_generator.generate_7, discriminator.discriminate_7, generator_optimizer, discriminator_optimizer, dataroot, device) # torch.save(generator.state_dict(), f'./split_generator') # torch.save(discriminator.state_dict(), f'./split_discriminator') # torch.save(inference_generator.state_dict(), f'./split_inference_generator') # generator.load_state_dict(torch.load('./generator')) # discriminator.load_state_dict(torch.load('./discriminator')) # interpolation_step( # 48, 400, test_noise, # generator, inference_generator, discriminator, accumulate, # generator.generate_8, inference_generator.generate_8, # discriminator.discriminate_8, # generator_optimizer, discriminator_optimizer, # dataroot, device # ) # step( # 48, 400, test_noise, # generator, inference_generator, discriminator, accumulate, # generator.generate_9, inference_generator.generate_9, # discriminator.discriminate_9, # generator_optimizer, discriminator_optimizer, # dataroot, device # ) torch.save(generator.state_dict(), f'./split_90_generator_final') torch.save(discriminator.state_dict(), f'./split_90_discriminator_final') torch.save(inference_generator.state_dict(), f'./split_90_inference_generator_final')
def main(args): #transformer transform = transforms.Compose([ transforms.Resize(64), transforms.ToTensor(), transforms.Normalize(mean=[.5, .5, .5], std=[.5, .5, .5]) ]) #dateset anime = AnimeData(args.tags, args.imgs, transform=transform) dataloder = DataLoader(anime, batch_size=args.bs, shuffle=True) #model gen = Generator(args.noise, momentum=args.momentum) dis = Discriminator(momentum=args.momentum) #criterion criterion = nn.BCELoss() if torch.cuda.is_available(): gen = gen.cuda() dis = dis.cuda() criterion = criterion.cuda() #optimizer optimizer_gen = optim.Adam(gen.parameters(), lr=args.lr_g, betas=args.betas) optimizer_dis = optim.Adam(dis.parameters(), lr=args.lr_d, betas=args.betas) loss_history_d = [] loss_history_g = [] out_history_true_d = [] out_history_fake_d = [] out_history_fake_g = [] for epoch in range(args.epochs): print('----------------start epoch %d ---------------' % epoch) step = 0 for data in dataloder: step += 1 start = time.time() img = Variable(data) noise = Variable(torch.randn(img.shape[0], args.noise)) labels_true = Variable(torch.ones(img.shape[0], 1)) labels_fake = Variable(torch.zeros(img.shape[0], 1)) if args.label_smoothing: labels_true = labels_true - torch.rand(img.shape[0], 1) * 0.1 labels_fake = labels_fake + torch.rand(img.shape[0], 1) * 0.1 #train on GPU if torch.cuda.is_available(): img = img.cuda() noise = noise.cuda() labels_true = labels_true.cuda() labels_fake = labels_fake.cuda() #train D out_true_d = dis(img) out_fake_d = dis(gen(noise)) out_history_true_d.append(torch.mean(out_true_d).item()) out_history_fake_d.append(torch.mean(out_fake_d).item()) #d_loss_ture = -torch.mean(labels_true * torch.log(out_true_d) + (1. - labels_true) * torch.log(1. - out_true_d)) loss_true_d = criterion(out_true_d, labels_true) loss_fake_d = criterion(out_fake_d, labels_fake) loss_d = loss_true_d + loss_fake_d optimizer_dis.zero_grad() loss_d.backward() if args.check: print('>>>>>>>>>>check_d_grad<<<<<<<<<<') try: check_grad(dis, 'conv2.weight') except ValueError as e: print(e) show(loss_history_d, loss_history_g, out_history_true_d, out_history_fake_d, out_history_fake_g) torch.save(dis.state_dict(), os.path.join(os.getcwd(), args.d, 'bad.pth')) torch.save(gen.state_dict(), os.path.join(os.getcwd(), args.g, 'bad.pth')) return loss_history_d.append(loss_d.item()) optimizer_dis.step() #train G noise = Variable(torch.randn(img.shape[0], args.noise)) if torch.cuda.is_available(): noise = noise.cuda() out_fake_g = dis(gen(noise)) labels_fake = 1. - labels_fake out_history_fake_g.append(torch.mean(out_fake_g).item()) loss_g = criterion(out_fake_g, labels_fake) optimizer_gen.zero_grad() loss_g.backward() if args.check: print('>>>>>>>>>>check_g_grad<<<<<<<<<<') try: check_grad(gen, 'convTrans.weight') except ValueError as e: print(e) show(loss_history_d, loss_history_g, out_history_true_d, out_history_fake_d, out_history_fake_g) torch.save(dis.state_dict(), os.path.join(os.getcwd(), args.d, 'bad.pth')) torch.save(gen.state_dict(), os.path.join(os.getcwd(), args.g, 'bad.pth')) return loss_history_g.append(loss_g.item()) optimizer_gen.step() end = time.time() print( 'epoch: %d step: %d d_true: %.2f d_fake: %.2f g_fake: %.2f time: %.2f' % (epoch, step, out_history_true_d[-1], out_history_fake_d[-1], out_history_fake_g[-1], end - start)) #save model torch.save(dis.state_dict(), os.path.join(os.getcwd(), args.d, '{}.pth'.format(epoch))) torch.save(gen.state_dict(), os.path.join(os.getcwd(), args.g, '{}.pth'.format(epoch)))
def main(args): use_cuda = (len(args.gpuid) >= 1) print("{0} GPU(s) are available".format(cuda.device_count())) # Load dataset splits = ['train', 'valid'] if data.has_binary_files(args.data, splits): dataset = data.load_dataset(args.data, splits, args.src_lang, args.trg_lang, args.fixed_max_len) else: dataset = data.load_raw_text_dataset(args.data, splits, args.src_lang, args.trg_lang, args.fixed_max_len) if args.src_lang is None or args.trg_lang is None: # record inferred languages in args, so that it's saved in checkpoints args.src_lang, args.trg_lang = dataset.src, dataset.dst print('| [{}] dictionary: {} types'.format(dataset.src, len(dataset.src_dict))) print('| [{}] dictionary: {} types'.format(dataset.dst, len(dataset.dst_dict))) for split in splits: print('| {} {} {} examples'.format(args.data, split, len(dataset.splits[split]))) g_logging_meters = OrderedDict() g_logging_meters['train_loss'] = AverageMeter() g_logging_meters['valid_loss'] = AverageMeter() g_logging_meters['train_acc'] = AverageMeter() g_logging_meters['valid_acc'] = AverageMeter() g_logging_meters['bsz'] = AverageMeter() # sentences per batch d_logging_meters = OrderedDict() d_logging_meters['train_loss'] = AverageMeter() d_logging_meters['valid_loss'] = AverageMeter() d_logging_meters['train_acc'] = AverageMeter() d_logging_meters['valid_acc'] = AverageMeter() d_logging_meters['bsz'] = AverageMeter() # sentences per batch # Set model parameters args.encoder_embed_dim = 1000 args.encoder_layers = 2 # 4 args.encoder_dropout_out = 0 args.decoder_embed_dim = 1000 args.decoder_layers = 2 # 4 args.decoder_out_embed_dim = 1000 args.decoder_dropout_out = 0 args.bidirectional = False generator = LSTMModel(args, dataset.src_dict, dataset.dst_dict, use_cuda=use_cuda) print("Generator loaded successfully!") discriminator = Discriminator(args, dataset.src_dict, dataset.dst_dict, use_cuda=use_cuda) print("Discriminator loaded successfully!") g_model_path = 'checkpoints/zhenwarm/generator.pt' assert os.path.exists(g_model_path) # generator = LSTMModel(args, dataset.src_dict, dataset.dst_dict, use_cuda=use_cuda) model_dict = generator.state_dict() model = torch.load(g_model_path) pretrained_dict = model.state_dict() # 1. filter out unnecessary keys pretrained_dict = { k: v for k, v in pretrained_dict.items() if k in model_dict } # 2. overwrite entries in the existing state dict model_dict.update(pretrained_dict) # 3. load the new state dict generator.load_state_dict(model_dict) print("pre-trained Generator loaded successfully!") # # Load discriminator model d_model_path = 'checkpoints/zhenwarm/discri.pt' assert os.path.exists(d_model_path) # generator = LSTMModel(args, dataset.src_dict, dataset.dst_dict, use_cuda=use_cuda) d_model_dict = discriminator.state_dict() d_model = torch.load(d_model_path) d_pretrained_dict = d_model.state_dict() # 1. filter out unnecessary keys d_pretrained_dict = { k: v for k, v in d_pretrained_dict.items() if k in d_model_dict } # 2. overwrite entries in the existing state dict d_model_dict.update(d_pretrained_dict) # 3. load the new state dict discriminator.load_state_dict(d_model_dict) print("pre-trained Discriminator loaded successfully!") if use_cuda: if torch.cuda.device_count() > 1: discriminator = torch.nn.DataParallel(discriminator).cuda() generator = torch.nn.DataParallel(generator).cuda() else: generator.cuda() discriminator.cuda() else: discriminator.cpu() generator.cpu() # adversarial training checkpoints saving path if not os.path.exists('checkpoints/myzhencli5'): os.makedirs('checkpoints/myzhencli5') checkpoints_path = 'checkpoints/myzhencli5/' # define loss function g_criterion = torch.nn.NLLLoss(ignore_index=dataset.dst_dict.pad(), reduction='sum') d_criterion = torch.nn.BCELoss() pg_criterion = PGLoss(ignore_index=dataset.dst_dict.pad(), size_average=True, reduce=True) # fix discriminator word embedding (as Wu et al. do) for p in discriminator.embed_src_tokens.parameters(): p.requires_grad = False for p in discriminator.embed_trg_tokens.parameters(): p.requires_grad = False # define optimizer g_optimizer = eval("torch.optim." + args.g_optimizer)(filter( lambda x: x.requires_grad, generator.parameters()), args.g_learning_rate) d_optimizer = eval("torch.optim." + args.d_optimizer)( filter(lambda x: x.requires_grad, discriminator.parameters()), args.d_learning_rate, momentum=args.momentum, nesterov=True) # start joint training best_dev_loss = math.inf num_update = 0 # main training loop for epoch_i in range(1, args.epochs + 1): logging.info("At {0}-th epoch.".format(epoch_i)) seed = args.seed + epoch_i torch.manual_seed(seed) max_positions_train = (args.fixed_max_len, args.fixed_max_len) # Initialize dataloader, starting at batch_offset trainloader = dataset.train_dataloader( 'train', max_tokens=args.max_tokens, max_sentences=args.joint_batch_size, max_positions=max_positions_train, # seed=seed, epoch=epoch_i, sample_without_replacement=args.sample_without_replacement, sort_by_source_size=(epoch_i <= args.curriculum), shard_id=args.distributed_rank, num_shards=args.distributed_world_size, ) # reset meters for key, val in g_logging_meters.items(): if val is not None: val.reset() for key, val in d_logging_meters.items(): if val is not None: val.reset() for i, sample in enumerate(trainloader): # set training mode generator.train() discriminator.train() update_learning_rate(num_update, 8e4, args.g_learning_rate, args.lr_shrink, g_optimizer) if use_cuda: # wrap input tensors in cuda tensors sample = utils.make_variable(sample, cuda=cuda) ## part I: use gradient policy method to train the generator # use policy gradient training when random.random() > 50% if random.random() >= 0.5: print("Policy Gradient Training") sys_out_batch = generator(sample) # 64 X 50 X 6632 out_batch = sys_out_batch.contiguous().view( -1, sys_out_batch.size(-1)) # (64 * 50) X 6632 _, prediction = out_batch.topk(1) prediction = prediction.squeeze(1) # 64*50 = 3200 prediction = torch.reshape( prediction, sample['net_input']['src_tokens'].shape) # 64 X 50 with torch.no_grad(): reward = discriminator(sample['net_input']['src_tokens'], prediction) # 64 X 1 train_trg_batch = sample['target'] # 64 x 50 pg_loss = pg_criterion(sys_out_batch, train_trg_batch, reward, use_cuda) sample_size = sample['target'].size( 0) if args.sentence_avg else sample['ntokens'] # 64 logging_loss = pg_loss / math.log(2) g_logging_meters['train_loss'].update(logging_loss.item(), sample_size) logging.debug( f"G policy gradient loss at batch {i}: {pg_loss.item():.3f}, lr={g_optimizer.param_groups[0]['lr']}" ) g_optimizer.zero_grad() pg_loss.backward() torch.nn.utils.clip_grad_norm_(generator.parameters(), args.clip_norm) g_optimizer.step() else: # MLE training print("MLE Training") sys_out_batch = generator(sample) out_batch = sys_out_batch.contiguous().view( -1, sys_out_batch.size(-1)) # (64 X 50) X 6632 train_trg_batch = sample['target'].view(-1) # 64*50 = 3200 loss = g_criterion(out_batch, train_trg_batch) sample_size = sample['target'].size( 0) if args.sentence_avg else sample['ntokens'] nsentences = sample['target'].size(0) logging_loss = loss.data / sample_size / math.log(2) g_logging_meters['bsz'].update(nsentences) g_logging_meters['train_loss'].update(logging_loss, sample_size) logging.debug( f"G MLE loss at batch {i}: {g_logging_meters['train_loss'].avg:.3f}, lr={g_optimizer.param_groups[0]['lr']}" ) g_optimizer.zero_grad() loss.backward() # all-reduce grads and rescale by grad_denom for p in generator.parameters(): if p.requires_grad: p.grad.data.div_(sample_size) torch.nn.utils.clip_grad_norm_(generator.parameters(), args.clip_norm) g_optimizer.step() num_update += 1 # part II: train the discriminator if num_update % 5 == 0: bsz = sample['target'].size(0) # batch_size = 64 src_sentence = sample['net_input'][ 'src_tokens'] # 64 x max-len i.e 64 X 50 # now train with machine translation output i.e generator output true_sentence = sample['target'].view(-1) # 64*50 = 3200 true_labels = Variable( torch.ones( sample['target'].size(0)).float()) # 64 length vector with torch.no_grad(): sys_out_batch = generator(sample) # 64 X 50 X 6632 out_batch = sys_out_batch.contiguous().view( -1, sys_out_batch.size(-1)) # (64 X 50) X 6632 _, prediction = out_batch.topk(1) prediction = prediction.squeeze(1) # 64 * 50 = 6632 fake_labels = Variable( torch.zeros( sample['target'].size(0)).float()) # 64 length vector fake_sentence = torch.reshape(prediction, src_sentence.shape) # 64 X 50 true_sentence = torch.reshape(true_sentence, src_sentence.shape) if use_cuda: fake_labels = fake_labels.cuda() true_labels = true_labels.cuda() # fake_disc_out = discriminator(src_sentence, fake_sentence) # 64 X 1 # true_disc_out = discriminator(src_sentence, true_sentence) # # fake_d_loss = d_criterion(fake_disc_out.squeeze(1), fake_labels) # true_d_loss = d_criterion(true_disc_out.squeeze(1), true_labels) # # fake_acc = torch.sum(torch.round(fake_disc_out).squeeze(1) == fake_labels).float() / len(fake_labels) # true_acc = torch.sum(torch.round(true_disc_out).squeeze(1) == true_labels).float() / len(true_labels) # acc = (fake_acc + true_acc) / 2 # # d_loss = fake_d_loss + true_d_loss if random.random() > 0.5: fake_disc_out = discriminator(src_sentence, fake_sentence) fake_d_loss = d_criterion(fake_disc_out.squeeze(1), fake_labels) fake_acc = torch.sum( torch.round(fake_disc_out).squeeze(1) == fake_labels).float() / len(fake_labels) d_loss = fake_d_loss acc = fake_acc else: true_disc_out = discriminator(src_sentence, true_sentence) true_d_loss = d_criterion(true_disc_out.squeeze(1), true_labels) true_acc = torch.sum( torch.round(true_disc_out).squeeze(1) == true_labels).float() / len(true_labels) d_loss = true_d_loss acc = true_acc d_logging_meters['train_acc'].update(acc) d_logging_meters['train_loss'].update(d_loss) logging.debug( f"D training loss {d_logging_meters['train_loss'].avg:.3f}, acc {d_logging_meters['train_acc'].avg:.3f} at batch {i}" ) d_optimizer.zero_grad() d_loss.backward() d_optimizer.step() if num_update % 10000 == 0: # validation # set validation mode generator.eval() discriminator.eval() # Initialize dataloader max_positions_valid = (args.fixed_max_len, args.fixed_max_len) valloader = dataset.eval_dataloader( 'valid', max_tokens=args.max_tokens, max_sentences=args.joint_batch_size, max_positions=max_positions_valid, skip_invalid_size_inputs_valid_test=True, descending= True, # largest batch first to warm the caching allocator shard_id=args.distributed_rank, num_shards=args.distributed_world_size, ) # reset meters for key, val in g_logging_meters.items(): if val is not None: val.reset() for key, val in d_logging_meters.items(): if val is not None: val.reset() for i, sample in enumerate(valloader): with torch.no_grad(): if use_cuda: # wrap input tensors in cuda tensors sample = utils.make_variable(sample, cuda=cuda) # generator validation sys_out_batch = generator(sample) out_batch = sys_out_batch.contiguous().view( -1, sys_out_batch.size(-1)) # (64 X 50) X 6632 dev_trg_batch = sample['target'].view( -1) # 64*50 = 3200 loss = g_criterion(out_batch, dev_trg_batch) sample_size = sample['target'].size( 0) if args.sentence_avg else sample['ntokens'] loss = loss / sample_size / math.log(2) g_logging_meters['valid_loss'].update( loss, sample_size) logging.debug( f"G dev loss at batch {i}: {g_logging_meters['valid_loss'].avg:.3f}" ) # discriminator validation bsz = sample['target'].size(0) src_sentence = sample['net_input']['src_tokens'] # train with half human-translation and half machine translation true_sentence = sample['target'] true_labels = Variable( torch.ones(sample['target'].size(0)).float()) with torch.no_grad(): sys_out_batch = generator(sample) out_batch = sys_out_batch.contiguous().view( -1, sys_out_batch.size(-1)) # (64 X 50) X 6632 _, prediction = out_batch.topk(1) prediction = prediction.squeeze(1) # 64 * 50 = 6632 fake_labels = Variable( torch.zeros(sample['target'].size(0)).float()) fake_sentence = torch.reshape( prediction, src_sentence.shape) # 64 X 50 true_sentence = torch.reshape(true_sentence, src_sentence.shape) if use_cuda: fake_labels = fake_labels.cuda() true_labels = true_labels.cuda() fake_disc_out = discriminator(src_sentence, fake_sentence) # 64 X 1 true_disc_out = discriminator(src_sentence, true_sentence) fake_d_loss = d_criterion(fake_disc_out.squeeze(1), fake_labels) true_d_loss = d_criterion(true_disc_out.squeeze(1), true_labels) d_loss = fake_d_loss + true_d_loss fake_acc = torch.sum( torch.round(fake_disc_out).squeeze(1) == fake_labels).float() / len(fake_labels) true_acc = torch.sum( torch.round(true_disc_out).squeeze(1) == true_labels).float() / len(true_labels) acc = (fake_acc + true_acc) / 2 d_logging_meters['valid_acc'].update(acc) d_logging_meters['valid_loss'].update(d_loss) logging.debug( f"D dev loss {d_logging_meters['valid_loss'].avg:.3f}, acc {d_logging_meters['valid_acc'].avg:.3f} at batch {i}" ) # torch.save(discriminator, # open(checkpoints_path + f"numupdate_{num_update/10000}k.discri_{d_logging_meters['valid_loss'].avg:.3f}.pt",'wb'), pickle_module=dill) # if d_logging_meters['valid_loss'].avg < best_dev_loss: # best_dev_loss = d_logging_meters['valid_loss'].avg # torch.save(discriminator, open(checkpoints_path + "best_dmodel.pt", 'wb'), pickle_module=dill) torch.save( generator, open( checkpoints_path + f"numupdate_{num_update/10000}k.joint_{g_logging_meters['valid_loss'].avg:.3f}.pt", 'wb'), pickle_module=dill)
class GAN_CLS(object): def __init__(self, args, data_loader, SUPERVISED=True): """ args : Arguments data_loader = An instance of class DataLoader for loading our dataset in batches """ self.data_loader = data_loader self.num_epochs = args.num_epochs self.batch_size = args.batch_size self.log_step = args.log_step self.sample_step = args.sample_step self.log_dir = args.log_dir self.checkpoint_dir = args.checkpoint_dir self.sample_dir = args.sample_dir self.final_model = args.final_model self.model_save_step = args.model_save_step #self.dataset = args.dataset #self.model_name = args.model_name self.img_size = args.img_size self.z_dim = args.z_dim self.text_embed_dim = args.text_embed_dim self.text_reduced_dim = args.text_reduced_dim self.learning_rate = args.learning_rate self.beta1 = args.beta1 self.beta2 = args.beta2 self.l1_coeff = args.l1_coeff self.resume_epoch = args.resume_epoch self.resume_idx = args.resume_idx self.SUPERVISED = SUPERVISED # Logger setting log_name = datetime.datetime.now().strftime('%Y-%m-%d') + '.log' self.logger = logging.getLogger('__name__') self.logger.setLevel(logging.INFO) self.formatter = logging.Formatter( '%(asctime)s:%(levelname)s:%(message)s') self.file_handler = logging.FileHandler( os.path.join(self.log_dir, log_name)) self.file_handler.setFormatter(self.formatter) self.logger.addHandler(self.file_handler) self.build_model() def smooth_label(self, tensor, offset): return tensor + offset def dump_imgs(images_Array, name): with open('{}.pickle'.format(name), 'wb') as file: dump(images_Array, file) def build_model(self): """ A function of defining following instances : ----- Generator ----- Discriminator ----- Optimizer for Generator ----- Optimizer for Discriminator ----- Defining Loss functions """ # ---------------------------------------------------------------------# # 1. Network Initialization # # ---------------------------------------------------------------------# self.gen = Generator(batch_size=self.batch_size, img_size=self.img_size, z_dim=self.z_dim, text_embed_dim=self.text_embed_dim, text_reduced_dim=self.text_reduced_dim) self.disc = Discriminator(batch_size=self.batch_size, img_size=self.img_size, text_embed_dim=self.text_embed_dim, text_reduced_dim=self.text_reduced_dim) self.gen_optim = optim.Adam(self.gen.parameters(), lr=self.learning_rate, betas=(self.beta1, self.beta2)) self.disc_optim = optim.Adam(self.disc.parameters(), lr=self.learning_rate, betas=(self.beta1, self.beta2)) self.cls_gan_optim = optim.Adam(itertools.chain( self.gen.parameters(), self.disc.parameters()), lr=self.learning_rate, betas=(self.beta1, self.beta2)) print('------------- Generator Model Info ---------------') self.print_network(self.gen, 'G') print('------------------------------------------------') print('------------- Discriminator Model Info ---------------') self.print_network(self.disc, 'D') print('------------------------------------------------') self.criterion = nn.BCELoss().cuda() # self.CE_loss = nn.CrossEntropyLoss().cuda() # self.MSE_loss = nn.MSELoss().cuda() self.gen.train() self.disc.train() def print_network(self, model, name): """ A function for printing total number of model parameters """ num_params = 0 for p in model.parameters(): num_params += p.numel() print(model) print(name) print("Total number of parameters: {}".format(num_params)) def load_checkpoints(self, resume_epoch, idx): """Restore the trained generator and discriminator.""" print('Loading the trained models from epoch {} and iteration {}...'. format(resume_epoch, idx)) G_path = os.path.join(self.checkpoint_dir, '{}-{}-G.ckpt'.format(resume_epoch, idx)) D_path = os.path.join(self.checkpoint_dir, '{}-{}-D.ckpt'.format(resume_epoch, idx)) self.gen.load_state_dict( torch.load(G_path, map_location=lambda storage, loc: storage)) self.disc.load_state_dict( torch.load(D_path, map_location=lambda storage, loc: storage)) def train_model(self): data_loader = self.data_loader start_epoch = 0 if self.resume_epoch >= 0: start_epoch = self.resume_epoch self.load_checkpoints(self.resume_epoch, self.resume_idx) print('--------------- Model Training Started ---------------') start_time = time.time() for epoch in range(start_epoch, self.num_epochs): print("Epoch: {}".format(epoch + 1)) for idx, batch in enumerate(data_loader): print("Index: {}".format(idx + 1), end="\t") true_imgs = batch['true_imgs'] true_embed = batch['true_embds'] false_imgs = batch['false_imgs'] real_labels = torch.ones(true_imgs.size(0)) fake_labels = torch.zeros(true_imgs.size(0)) smooth_real_labels = torch.FloatTensor( self.smooth_label(real_labels.numpy(), -0.1)) true_imgs = Variable(true_imgs.float()).cuda() true_embed = Variable(true_embed.float()).cuda() false_imgs = Variable(false_imgs.float()).cuda() real_labels = Variable(real_labels).cuda() smooth_real_labels = Variable(smooth_real_labels).cuda() fake_labels = Variable(fake_labels).cuda() # ---------------------------------------------------------------# # 2. Training the generator # # ---------------------------------------------------------------# self.gen.zero_grad() z = Variable(torch.randn(true_imgs.size(0), self.z_dim)).cuda() fake_imgs = self.gen.forward(true_embed, z) fake_out, fake_logit = self.disc.forward(fake_imgs, true_embed) fake_out = Variable(fake_out.data, requires_grad=True).cuda() true_out, true_logit = self.disc.forward(true_imgs, true_embed) true_out = Variable(true_out.data, requires_grad=True).cuda() g_sf = self.criterion(fake_out, real_labels) #g_img = self.l1_coeff * nn.L1Loss()(fake_imgs, true_imgs) gen_loss = g_sf gen_loss.backward() self.gen_optim.step() # ---------------------------------------------------------------# # 3. Training the discriminator # # ---------------------------------------------------------------# self.disc.zero_grad() false_out, false_logit = self.disc.forward( false_imgs, true_embed) false_out = Variable(false_out.data, requires_grad=True) sr = self.criterion(true_out, smooth_real_labels) sw = self.criterion(true_out, fake_labels) sf = self.criterion(false_out, smooth_real_labels) disc_loss = torch.log(sr) + (torch.log(1 - sw) + torch.log(1 - sf)) / 2 disc_loss.backward() self.disc_optim.step() self.cls_gan_optim.step() # Logging loss = {} loss['G_loss'] = gen_loss.item() loss['D_loss'] = disc_loss.item() # ---------------------------------------------------------------# # 4. Logging INFO into log_dir # # ---------------------------------------------------------------# log = "" if (idx + 1) % self.log_step == 0: end_time = time.time() - start_time end_time = datetime.timedelta(seconds=end_time) log = "Elapsed [{}], Epoch [{}/{}], Idx [{}]".format( end_time, epoch + 1, self.num_epochs, idx) for net, loss_value in loss.items(): log += "{}: {:.4f}".format(net, loss_value) self.logger.info(log) print(log) """ # ---------------------------------------------------------------# # 5. Saving generated images # # ---------------------------------------------------------------# if (idx + 1) % self.sample_step == 0: concat_imgs = torch.cat((true_imgs, fake_imgs), 0) # ?????????? concat_imgs = (concat_imgs + 1) / 2 # out.clamp_(0, 1) save_path = os.path.join(self.sample_dir, '{}-{}-images.jpg'.format(epoch, idx + 1)) # concat_imgs.cpu().detach().numpy() self.dump_imgs(concat_imgs.cpu().numpy(), save_path) #save_image(concat_imgs.data.cpu(), self.sample_dir, nrow=1, padding=0) print ('Saved real and fake images into {}...'.format(self.sample_dir)) """ # ---------------------------------------------------------------# # 6. Saving the checkpoints & final model # # ---------------------------------------------------------------# if (idx + 1) % self.model_save_step == 0: G_path = os.path.join( self.checkpoint_dir, '{}-{}-G.ckpt'.format(epoch, idx + 1)) D_path = os.path.join( self.checkpoint_dir, '{}-{}-D.ckpt'.format(epoch, idx + 1)) torch.save(self.gen.state_dict(), G_path) torch.save(self.disc.state_dict(), D_path) print('Saved model checkpoints into {}...\n'.format( self.checkpoint_dir)) print('--------------- Model Training Completed ---------------') # Saving final model into final_model directory G_path = os.path.join(self.final_model, '{}-G.pth'.format('final')) D_path = os.path.join(self.final_model, '{}-D.pth'.format('final')) torch.save(self.gen.state_dict(), G_path) torch.save(self.disc.state_dict(), D_path) print('Saved final model into {}...'.format(self.final_model))
class BigGAN(): """Big GAN""" def __init__(self, device, dataloader, num_classes, configs): self.device = device self.dataloader = dataloader self.num_classes = num_classes # model settings & hyperparams # self.total_steps = configs.total_steps self.epochs = configs.epochs self.d_iters = configs.d_iters self.g_iters = configs.g_iters self.batch_size = configs.batch_size self.imsize = configs.imsize self.nz = configs.nz self.ngf = configs.ngf self.ndf = configs.ndf self.g_lr = configs.g_lr self.d_lr = configs.d_lr self.beta1 = configs.beta1 self.beta2 = configs.beta2 # instance noise self.inst_noise_sigma = configs.inst_noise_sigma self.inst_noise_sigma_iters = configs.inst_noise_sigma_iters # model logging and saving self.log_step = configs.log_step self.save_epoch = configs.save_epoch self.model_path = configs.model_path self.sample_path = configs.sample_path # pretrained self.pretrained_model = configs.pretrained_model # building self.build_model() # archive of all losses self.ave_d_losses = [] self.ave_d_losses_real = [] self.ave_d_losses_fake = [] self.ave_g_losses = [] if self.pretrained_model: self.load_pretrained() def build_model(self): """Initiate Generator and Discriminator""" self.G = Generator(self.nz, self.ngf, self.num_classes).to(self.device) self.D = Discriminator(self.ndf, self.num_classes).to(self.device) self.g_optimizer = optim.Adam( filter(lambda p: p.requires_grad, self.G.parameters()), self.g_lr, [self.beta1, self.beta2]) self.d_optimizer = optim.Adam( filter(lambda p: p.requires_grad, self.D.parameters()), self.d_lr, [self.beta1, self.beta2]) print("Generator Parameters: ", parameters(self.G)) print(self.G) print("Discriminator Parameters: ", parameters(self.D)) print(self.D) print("Number of classes: ", self.num_classes) def load_pretrained(self): """Loading pretrained model""" checkpoint = torch.load( os.path.join(self.model_path, "{}_biggan.pth".format(self.pretrained_model))) # load models self.G.load_state_dict(checkpoint["g_state_dict"]) self.D.load_state_dict(checkpoint["d_state_dict"]) # load optimizers self.g_optimizer.load_state_dict(checkpoint["g_optimizer"]) self.d_optimizer.load_state_dict(checkpoint["d_optimizer"]) # load losses self.ave_d_losses = checkpoint["ave_d_losses"] self.ave_d_losses_real = checkpoint["ave_d_losses_real"] self.ave_d_losses_fake = checkpoint["ave_d_losses_fake"] self.ave_g_losses = checkpoint["ave_g_losses"] print("Loading pretrained models (epoch: {})..!".format( self.pretrained_model)) def reset_grad(self): """Reset gradients""" self.g_optimizer.zero_grad() self.d_optimizer.zero_grad() def train(self): """Train model""" step_per_epoch = len(self.dataloader) epochs = self.epochs total_steps = epochs * step_per_epoch # fixed z and labels for sampling generator images fixed_z = tensor2var(torch.randn(self.batch_size, self.nz), device=self.device) fixed_labels = tensor2var(torch.from_numpy( np.tile(np.arange(self.num_classes), self.batch_size)).long(), device=self.device) print("Initiating Training") print("Epochs: {}, Total Steps: {}, Steps/Epoch: {}".format( epochs, total_steps, step_per_epoch)) if self.pretrained_model: start_epoch = self.pretrained_model else: start_epoch = 0 self.D.train() self.G.train() # Instance noise - make random noise mean (0) and std for injecting inst_noise_mean = torch.full( (self.batch_size, 3, self.imsize, self.imsize), 0).to(self.device) inst_noise_std = torch.full( (self.batch_size, 3, self.imsize, self.imsize), self.inst_noise_sigma).to(self.device) # total time start_time = time.time() for epoch in range(start_epoch, epochs): # local losses d_losses = [] d_losses_real = [] d_losses_fake = [] g_losses = [] data_iter = iter(self.dataloader) for step in range(step_per_epoch): # Instance noise std is linearly annealed from self.inst_noise_sigma to 0 thru self.inst_noise_sigma_iters inst_noise_sigma_curr = 0 if step > self.inst_noise_sigma_iters else ( 1 - step / self.inst_noise_sigma_iters) * self.inst_noise_sigma inst_noise_std.fill_(inst_noise_sigma_curr) # get real images real_images, real_labels = next(data_iter) real_images = real_images.to(self.device) real_labels = real_labels.to(self.device) # ================== TRAIN DISCRIMINATOR ================== # for _ in range(self.d_iters): self.reset_grad() # TRAIN REAL # creating instance noise inst_noise = torch.normal(mean=inst_noise_mean, std=inst_noise_std).to( self.device) # adding noise to real images d_real = self.D(real_images + inst_noise, real_labels) d_loss_real = loss_hinge_dis_real(d_real) d_loss_real.backward() # delete loss if (step + 1) % self.log_step != 0: del d_real, d_loss_real # TRAIN FAKE # create fake images using latent vector z = tensor2var(torch.randn(real_images.size(0), self.nz), device=self.device) fake_images = self.G(z, real_labels) # creating instance noise inst_noise = torch.normal(mean=inst_noise_mean, std=inst_noise_std).to( self.device) # adding noise to fake images # detach fake_images tensor from graph d_fake = self.D(fake_images.detach() + inst_noise, real_labels) d_loss_fake = loss_hinge_dis_fake(d_fake) d_loss_fake.backward() # delete loss, output del fake_images if (step + 1) % self.log_step != 0: del d_fake, d_loss_fake # optimize D self.d_optimizer.step() # ================== TRAIN GENERATOR ================== # for _ in range(self.g_iters): self.reset_grad() # create new latent vector z = tensor2var(torch.randn(real_images.size(0), self.nz), device=self.device) # generate fake images inst_noise = torch.normal(mean=inst_noise_mean, std=inst_noise_std).to( self.device) fake_images = self.G(z, real_labels) g_fake = self.D(fake_images + inst_noise, real_labels) # compute hinge loss for G g_loss = loss_hinge_gen(g_fake) g_loss.backward() del fake_images if (step + 1) % self.log_step != 0: del g_fake, g_loss # optimize G self.g_optimizer.step() # logging step progression if (step + 1) % self.log_step == 0: d_loss = d_loss_real + d_loss_fake # logging losses d_losses.append(d_loss.item()) d_losses_real.append(d_loss_real.item()) d_losses_fake.append(d_loss_fake.item()) g_losses.append(g_loss.item()) # print out elapsed = time.time() - start_time elapsed = str(datetime.timedelta(seconds=elapsed)) print( "Elapsed [{}], Epoch: [{}/{}], Step [{}/{}], g_loss: {:.4f}, d_loss: {:.4f}," " d_loss_real: {:.4f}, d_loss_fake: {:.4f}".format( elapsed, (epoch + 1), epochs, (step + 1), step_per_epoch, g_loss, d_loss, d_loss_real, d_loss_fake)) del d_real, d_loss_real, d_fake, d_loss_fake, g_fake, g_loss # logging average losses over epoch self.ave_d_losses.append(mean(d_losses)) self.ave_d_losses_real.append(mean(d_losses_real)) self.ave_d_losses_fake.append(mean(d_losses_fake)) self.ave_g_losses.append(mean(g_losses)) # epoch update print( "Elapsed [{}], Epoch: [{}/{}], ave_g_loss: {:.4f}, ave_d_loss: {:.4f}," " ave_d_loss_real: {:.4f}, ave_d_loss_fake: {:.4f},".format( elapsed, epoch + 1, epochs, self.ave_g_losses[epoch], self.ave_d_losses[epoch], self.ave_d_losses_real[epoch], self.ave_d_losses_fake[epoch])) # sample images every epoch fake_images = self.G(fixed_z, fixed_labels) fake_images = denorm(fake_images.data) save_image( fake_images, os.path.join(self.sample_path, "Epoch {}.png".format(epoch + 1))) # save model if (epoch + 1) % self.save_epoch == 0: torch.save( { "g_state_dict": self.G.state_dict(), "d_state_dict": self.D.state_dict(), "g_optimizer": self.g_optimizer.state_dict(), "d_optimizer": self.d_optimizer.state_dict(), "ave_d_losses": self.ave_d_losses, "ave_d_losses_real": self.ave_d_losses_real, "ave_d_losses_fake": self.ave_d_losses_fake, "ave_g_losses": self.ave_g_losses }, os.path.join(self.model_path, "{}_biggan.pth".format(epoch + 1))) print("Saving models (epoch {})..!".format(epoch + 1)) def plot(self): plt.plot(self.ave_d_losses) plt.plot(self.ave_d_losses_real) plt.plot(self.ave_d_losses_fake) plt.plot(self.ave_g_losses) plt.legend(["d loss", "d real", "d fake", "g loss"], loc="upper left") plt.show()
def main(args): # log hyperparameter print(args) # select device args.cuda = not args.no_cuda and torch.cuda.is_available() device = torch.device("cuda: 0" if args.cuda else "cpu") # set random seed np.random.seed(args.seed) torch.manual_seed(args.seed) # data loader transform = transforms.Compose([ utils.Normalize(), utils.ToTensor() ]) train_dataset = TVDataset( root=args.root, sub_size=args.block_size, volume_list=args.volume_train_list, max_k=args.training_step, train=True, transform=transform ) test_dataset = TVDataset( root=args.root, sub_size=args.block_size, volume_list=args.volume_test_list, max_k=args.training_step, train=False, transform=transform ) kwargs = {"num_workers": 4, "pin_memory": True} if args.cuda else {} train_loader = DataLoader(train_dataset, batch_size=args.batch_size, shuffle=True, **kwargs) test_loader = DataLoader(test_dataset, batch_size=args.batch_size, shuffle=False, **kwargs) # model def generator_weights_init(m): if isinstance(m, nn.Conv3d): nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu') if m.bias is not None: nn.init.zeros_(m.bias) def discriminator_weights_init(m): if isinstance(m, nn.Conv3d): nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='leaky_relu') if m.bias is not None: nn.init.zeros_(m.bias) g_model = Generator(args.upsample_mode, args.forward, args.backward, args.gen_sn, args.residual) g_model.apply(generator_weights_init) if args.data_parallel and torch.cuda.device_count() > 1: g_model = nn.DataParallel(g_model) g_model.to(device) if args.gan_loss != "none": d_model = Discriminator(args.dis_sn) d_model.apply(discriminator_weights_init) # if args.dis_sn: # d_model = add_sn(d_model) if args.data_parallel and torch.cuda.device_count() > 1: d_model = nn.DataParallel(d_model) d_model.to(device) mse_loss = nn.MSELoss() adversarial_loss = nn.MSELoss() train_losses, test_losses = [], [] d_losses, g_losses = [], [] # optimizer g_optimizer = optim.Adam(g_model.parameters(), lr=args.lr, betas=(args.beta1, args.beta2)) if args.gan_loss != "none": d_optimizer = optim.Adam(d_model.parameters(), lr=args.d_lr, betas=(args.beta1, args.beta2)) Tensor = torch.cuda.FloatTensor if args.cuda else torch.FloatTensor # load checkpoint if args.resume: if os.path.isfile(args.resume): print("=> loading checkpoint {}".format(args.resume)) checkpoint = torch.load(args.resume) args.start_epoch = checkpoint["epoch"] g_model.load_state_dict(checkpoint["g_model_state_dict"]) # g_optimizer.load_state_dict(checkpoint["g_optimizer_state_dict"]) if args.gan_loss != "none": d_model.load_state_dict(checkpoint["d_model_state_dict"]) # d_optimizer.load_state_dict(checkpoint["d_optimizer_state_dict"]) d_losses = checkpoint["d_losses"] g_losses = checkpoint["g_losses"] train_losses = checkpoint["train_losses"] test_losses = checkpoint["test_losses"] print("=> load chekcpoint {} (epoch {})" .format(args.resume, checkpoint["epoch"])) # main loop for epoch in tqdm(range(args.start_epoch, args.epochs)): # training.. g_model.train() if args.gan_loss != "none": d_model.train() train_loss = 0. volume_loss_part = np.zeros(args.training_step) for i, sample in enumerate(train_loader): params = list(g_model.named_parameters()) # pdb.set_trace() # params[0][1].register_hook(lambda g: print("{}.grad: {}".format(params[0][0], g))) # adversarial ground truths real_label = Variable(Tensor(sample["v_i"].shape[0], sample["v_i"].shape[1], 1, 1, 1, 1).fill_(1.0), requires_grad=False) fake_label = Variable(Tensor(sample["v_i"].shape[0], sample["v_i"].shape[1], 1, 1, 1, 1).fill_(0.0), requires_grad=False) v_f = sample["v_f"].to(device) v_b = sample["v_b"].to(device) v_i = sample["v_i"].to(device) g_optimizer.zero_grad() fake_volumes = g_model(v_f, v_b, args.training_step, args.wo_ori_volume, args.norm) # adversarial loss # update discriminator if args.gan_loss != "none": avg_d_loss = 0. avg_d_loss_real = 0. avg_d_loss_fake = 0. for k in range(args.n_d): d_optimizer.zero_grad() decisions = d_model(v_i) d_loss_real = adversarial_loss(decisions, real_label) fake_decisions = d_model(fake_volumes.detach()) d_loss_fake = adversarial_loss(fake_decisions, fake_label) d_loss = d_loss_real + d_loss_fake d_loss.backward() avg_d_loss += d_loss.item() / args.n_d avg_d_loss_real += d_loss_real / args.n_d avg_d_loss_fake += d_loss_fake / args.n_d d_optimizer.step() # update generator if args.gan_loss != "none": avg_g_loss = 0. avg_loss = 0. for k in range(args.n_g): loss = 0. g_optimizer.zero_grad() # adversarial loss if args.gan_loss != "none": fake_decisions = d_model(fake_volumes) g_loss = args.gan_loss_weight * adversarial_loss(fake_decisions, real_label) loss += g_loss avg_g_loss += g_loss.item() / args.n_g # volume loss if args.volume_loss: volume_loss = args.volume_loss_weight * mse_loss(v_i, fake_volumes) for j in range(v_i.shape[1]): volume_loss_part[j] += mse_loss(v_i[:, j, :], fake_volumes[:, j, :]) / args.n_g / args.log_every loss += volume_loss # feature loss if args.feature_loss: feat_real = d_model.extract_features(v_i) feat_fake = d_model.extract_features(fake_volumes) for m in range(len(feat_real)): loss += args.feature_loss_weight / len(feat_real) * mse_loss(feat_real[m], feat_fake[m]) avg_loss += loss / args.n_g loss.backward() g_optimizer.step() train_loss += avg_loss # log training status subEpoch = (i + 1) // args.log_every if (i+1) % args.log_every == 0: print("Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}".format( epoch, (i+1) * args.batch_size, len(train_loader.dataset), 100. * (i+1) / len(train_loader), avg_loss )) print("Volume Loss: ") for j in range(volume_loss_part.shape[0]): print("\tintermediate {}: {:.6f}".format( j+1, volume_loss_part[j] )) if args.gan_loss != "none": print("DLossReal: {:.6f} DLossFake: {:.6f} DLoss: {:.6f}, GLoss: {:.6f}".format( avg_d_loss_real, avg_d_loss_fake, avg_d_loss, avg_g_loss )) d_losses.append(avg_d_loss) g_losses.append(avg_g_loss) # train_losses.append(avg_loss) train_losses.append(train_loss.item() / args.log_every) print("====> SubEpoch: {} Average loss: {:.6f} Time {}".format( subEpoch, train_loss.item() / args.log_every, time.asctime(time.localtime(time.time())) )) train_loss = 0. volume_loss_part = np.zeros(args.training_step) # testing... if (i + 1) % args.test_every == 0: g_model.eval() if args.gan_loss != "none": d_model.eval() test_loss = 0. with torch.no_grad(): for i, sample in enumerate(test_loader): v_f = sample["v_f"].to(device) v_b = sample["v_b"].to(device) v_i = sample["v_i"].to(device) fake_volumes = g_model(v_f, v_b, args.training_step, args.wo_ori_volume, args.norm) test_loss += args.volume_loss_weight * mse_loss(v_i, fake_volumes).item() test_losses.append(test_loss * args.batch_size / len(test_loader.dataset)) print("====> SubEpoch: {} Test set loss {:4f} Time {}".format( subEpoch, test_losses[-1], time.asctime(time.localtime(time.time())) )) # saving... if (i+1) % args.check_every == 0: print("=> saving checkpoint at epoch {}".format(epoch)) if args.gan_loss != "none": torch.save({"epoch": epoch + 1, "g_model_state_dict": g_model.state_dict(), "g_optimizer_state_dict": g_optimizer.state_dict(), "d_model_state_dict": d_model.state_dict(), "d_optimizer_state_dict": d_optimizer.state_dict(), "d_losses": d_losses, "g_losses": g_losses, "train_losses": train_losses, "test_losses": test_losses}, os.path.join(args.save_dir, "model_" + str(epoch) + "_" + str(subEpoch) + "_" + "pth.tar") ) else: torch.save({"epoch": epoch + 1, "g_model_state_dict": g_model.state_dict(), "g_optimizer_state_dict": g_optimizer.state_dict(), "train_losses": train_losses, "test_losses": test_losses}, os.path.join(args.save_dir, "model_" + str(epoch) + "_" + str(subEpoch) + "_" + "pth.tar") ) torch.save(g_model.state_dict(), os.path.join(args.save_dir, "model_" + str(epoch) + "_" + str(subEpoch) + ".pth")) num_subEpoch = len(train_loader) // args.log_every print("====> Epoch: {} Average loss: {:.6f} Time {}".format( epoch, np.array(train_losses[-num_subEpoch:]).mean(), time.asctime(time.localtime(time.time())) ))
np.save(os.path.join(save_dir, 'd_losses.npy'), d_losses) np.save(os.path.join(save_dir, 'g_losses.npy'), g_losses) np.save(os.path.join(save_dir, 'fake_scores.npy'), fake_scores) np.save(os.path.join(save_dir, 'real_scores.npy'), real_scores) plt.figure() pylab.xlim(0, num_epochs + 1) plt.plot(range(1, num_epochs + 1), d_losses, label='d loss') plt.plot(range(1, num_epochs + 1), g_losses, label='g loss') plt.legend() plt.savefig(os.path.join(save_dir, 'loss.pdf')) plt.close() plt.figure() pylab.xlim(0, num_epochs + 1) pylab.ylim(0, 1) plt.plot(range(1, num_epochs + 1), fake_scores, label='fake score') plt.plot(range(1, num_epochs + 1), real_scores, label='real score') plt.legend() plt.savefig(os.path.join(save_dir, 'accuracy.pdf')) plt.close() # Save model at checkpoints if (epoch + 1) % 50 == 0: torch.save(G.state_dict(), os.path.join(save_dir, 'G--{}.ckpt'.format(epoch+1))) torch.save(D.state_dict(), os.path.join(save_dir, 'D--{}.ckpt'.format(epoch+1))) # Save the model checkpoints torch.save(G.state_dict(), 'G.ckpt') torch.save(D.state_dict(), 'D.ckpt')
def main(): random.seed(SEED) np.random.seed(SEED) calc_bleu([1, 10, 12]) exit() # Build up dataset s_train, s_test = load_from_big_file('../data/train_data_obama.txt') # idx_to_word: List of id to word # word_to_idx: Dictionary mapping word to id idx_to_word, word_to_idx = fetch_vocab(s_train, s_train, s_test) # TODO: 1. Prepare data for attention model # input_seq, target_seq = prepare_data(DATA_GERMAN, DATA_ENGLISH, word_to_idx) global VOCAB_SIZE VOCAB_SIZE = len(idx_to_word) save_vocab(CHECKPOINT_PATH + 'metadata.data', idx_to_word, word_to_idx, VOCAB_SIZE, g_emb_dim, g_hidden_dim, g_sequence_len) print('VOCAB SIZE:', VOCAB_SIZE) # Define Networks generator = Generator(VOCAB_SIZE, g_emb_dim, g_hidden_dim, g_sequence_len, BATCH_SIZE, opt.cuda) discriminator = Discriminator(d_num_class, VOCAB_SIZE, d_emb_dim, d_filter_sizes, d_num_filters, d_dropout) target_lstm = TargetLSTM(VOCAB_SIZE, g_emb_dim, g_hidden_dim, opt.cuda) if opt.cuda: generator = generator.cuda() discriminator = discriminator.cuda() target_lstm = target_lstm.cuda() # Generate toy data using target lstm print('Generating data ...') generate_real_data('../data/train_data_obama.txt', BATCH_SIZE, GENERATED_NUM, idx_to_word, word_to_idx, POSITIVE_FILE, TEST_FILE) # Create Test data iterator for testing test_iter = GenDataIter(TEST_FILE, BATCH_SIZE) # generate_samples(target_lstm, BATCH_SIZE, GENERATED_NUM, POSITIVE_FILE, idx_to_word) # Load data from file gen_data_iter = GenDataIter(POSITIVE_FILE, BATCH_SIZE) # Pretrain Generator using MLE # gen_criterion = nn.NLLLoss(size_average=False) gen_criterion = nn.CrossEntropyLoss() gen_optimizer = optim.Adam(generator.parameters()) if opt.cuda: gen_criterion = gen_criterion.cuda() print('Pretrain with MLE ...') for epoch in range(PRE_EPOCH_NUM): loss = train_epoch(generator, gen_data_iter, gen_criterion, gen_optimizer) print('Epoch [%d] Model Loss: %f' % (epoch, loss)) print('Training Output') test_predict(generator, test_iter, idx_to_word, train_mode=True) sys.stdout.flush() # TODO: 2. Flags to ensure dimension of model input is handled # generate_samples(generator, BATCH_SIZE, GENERATED_NUM, EVAL_FILE) """ eval_iter = GenDataIter(EVAL_FILE, BATCH_SIZE) print('Iterator Done') loss = eval_epoch(target_lstm, eval_iter, gen_criterion) print('Epoch [%d] True Loss: %f' % (epoch, loss)) """ print('OUTPUT AFTER PRE-TRAINING') test_predict(generator, test_iter, idx_to_word, train_mode=True) # Pretrain Discriminator dis_criterion = nn.NLLLoss(size_average=False) dis_optimizer = optim.Adam(discriminator.parameters()) if opt.cuda: dis_criterion = dis_criterion.cuda() print('Pretrain Discriminator ...') for epoch in range(3): generate_samples(generator, BATCH_SIZE, GENERATED_NUM, NEGATIVE_FILE) dis_data_iter = DisDataIter(POSITIVE_FILE, NEGATIVE_FILE, BATCH_SIZE) for _ in range(3): loss = train_epoch(discriminator, dis_data_iter, dis_criterion, dis_optimizer) print('Epoch [%d], loss: %f' % (epoch, loss)) sys.stdout.flush() # Adversarial Training rollout = Rollout(generator, 0.8) print('#####################################################') print('Start Adversarial Training...\n') gen_gan_loss = GANLoss() gen_gan_optm = optim.Adam(generator.parameters()) if opt.cuda: gen_gan_loss = gen_gan_loss.cuda() gen_criterion = nn.NLLLoss(size_average=False) if opt.cuda: gen_criterion = gen_criterion.cuda() dis_criterion = nn.NLLLoss(size_average=False) dis_optimizer = optim.Adam(discriminator.parameters()) if opt.cuda: dis_criterion = dis_criterion.cuda() real_iter = GenDataIter(POSITIVE_FILE, BATCH_SIZE) for total_batch in range(TOTAL_BATCH): ## Train the generator for one step for it in range(1): if real_iter.idx >= real_iter.data_num: real_iter.reset() inputs = real_iter.next()[0] inputs = inputs.cuda() samples = generator.sample(BATCH_SIZE, g_sequence_len, inputs) samples = samples.cpu() rewards = rollout.get_reward(samples, 16, discriminator) rewards = Variable(torch.Tensor(rewards)) if opt.cuda: rewards = torch.exp(rewards.cuda()).contiguous().view((-1, )) prob = generator.forward(inputs) mini_batch = prob.shape[0] prob = torch.reshape( prob, (prob.shape[0] * prob.shape[1], -1)) #prob.view(-1, g_emb_dim) targets = copy.deepcopy(inputs).contiguous().view((-1, )) loss = gen_gan_loss(prob, targets, rewards) gen_gan_optm.zero_grad() loss.backward() gen_gan_optm.step() """ samples = generator.sample(BATCH_SIZE, g_sequence_len) # construct the input to the genrator, add zeros before samples and delete the last column zeros = torch.zeros((BATCH_SIZE, 1)).type(torch.LongTensor) if samples.is_cuda: zeros = zeros.cuda() inputs = Variable(torch.cat([zeros, samples.data], dim = 1)[:, :-1].contiguous()) targets = Variable(samples.data).contiguous().view((-1,)) print('', inputs.shape, targets.shape) print(inputs, targets) # calculate the reward rewards = rollout.get_reward(samples, 16, discriminator) rewards = Variable(torch.Tensor(rewards)) if opt.cuda: rewards = torch.exp(rewards.cuda()).contiguous().view((-1,)) prob = generator.forward(inputs) mini_batch = prob.shape[0] prob = torch.reshape(prob, (prob.shape[0] * prob.shape[1], -1)) #prob.view(-1, g_emb_dim) loss = gen_gan_loss(prob, targets, rewards) gen_gan_optm.zero_grad() loss.backward() gen_gan_optm.step() """ print('Batch [%d] True Loss: %f' % (total_batch, loss)) if total_batch % 1 == 0 or total_batch == TOTAL_BATCH - 1: # generate_samples(generator, BATCH_SIZE, GENERATED_NUM, EVAL_FILE) # eval_iter = GenDataIter(EVAL_FILE, BATCH_SIZE) # loss = eval_epoch(target_lstm, eval_iter, gen_criterion) if len(prob.shape) > 2: prob = torch.reshape(prob, (prob.shape[0] * prob.shape[1], -1)) predictions = torch.max(prob, dim=1)[1] predictions = predictions.view(mini_batch, -1) for each_sen in list(predictions): print('Train Output:', generate_sentence_from_id(idx_to_word, each_sen)) test_predict(generator, test_iter, idx_to_word, train_mode=True) torch.save(generator.state_dict(), CHECKPOINT_PATH + 'generator.model') torch.save(discriminator.state_dict(), CHECKPOINT_PATH + 'discriminator.model') rollout.update_params() for _ in range(4): generate_samples(generator, BATCH_SIZE, GENERATED_NUM, NEGATIVE_FILE) dis_data_iter = DisDataIter(POSITIVE_FILE, NEGATIVE_FILE, BATCH_SIZE) for _ in range(2): loss = train_epoch(discriminator, dis_data_iter, dis_criterion, dis_optimizer)
'deblur_average_mse_loss'] += mse_loss.data.item() loss_values[ 'deblur_total_average_loss'] += total_loss.data.item() # ===================log======================== loss_values = { k: v / num_epochs for k, v in loss_values.items() } losses_per_epoch.append(loss_values) print('epoch [{}/{}], {}'.format(epoch + 1, num_epochs, loss_values)) #print('epoch [{}/{}], Deblurrer Total Average Loss: {:.4f}, ' + # 'Discrim Average Loss: {:.4f}, ' #.format(epoch + 1, num_epochs, total_loss.data, discrim_total_error.data)) except KeyboardInterrupt: torch.save(model.state_dict(), 'semantic_model_interrupt.pth') torch.save(discriminator.state_dict(), 'discrim_interrupt.pth') f = open("losses.txt", "w") f.write(str(losses_per_epoch)) f.close() sys.exit() break torch.save(model.state_dict(), 'semanticmodel.pth') torch.save(discriminator.state_dict(), 'discrim.pth') f = open("losses.txt", "w") f.write(str(losses_per_epoch)) f.close()
class GAN: def __init__(self, device, args): self.device = device self.args = args self.batch_size = args.batch_size self.generator_checkpoint_path = os.path.join(args.checkpoint_path, 'generator.pth') self.discriminator_checkpoint_path = os.path.join(args.checkpoint_path, 'discriminator.pth') if not os.path.isdir(args.checkpoint_path): os.mkdir(args.checkpoint_path) self.generator = Generator(args).to(self.device) self.discriminator = Discriminator(args).to(self.device) self.sequence_loss = SequenceLoss() self.reinforce_loss = ReinforceLoss() self.generator_optimizer = optim.Adam(self.generator.parameters(), lr=args.generator_lr) self.discriminator_optimizer = optim.Adam(self.discriminator.parameters(), lr=args.discriminator_lr) self.evaluator = Evaluator('val', self.device, args) self.cider = Cider(args) generator_dataset = CaptionDataset('train', args) self.generator_loader = DataLoader(generator_dataset, batch_size=self.batch_size, shuffle=True, num_workers=4) discriminator_dataset = DiscCaption('train', args) self.discriminator_loader = DataLoader(discriminator_dataset, batch_size=self.batch_size, shuffle=True, num_workers=4) def train(self): if self.args.load_generator: self.generator.load_state_dict(torch.load(self.generator_checkpoint_path)) else: self._pretrain_generator() if self.args.load_discriminator: self.discriminator.load_state_dict(torch.load(self.discriminator_checkpoint_path)) else: self._pretrain_discriminator() self._train_gan() def _pretrain_generator(self): iter = 0 for epoch in range(self.args.pretrain_generator_epochs): self.generator.train() for data in self.generator_loader: for name, item in data.items(): data[name] = item.to(self.device) self.generator.zero_grad() probs = self.generator(data['fc_feats'], data['att_feats'], data['att_masks'], data['labels']) loss = self.sequence_loss(probs, data['labels']) loss.backward() self.generator_optimizer.step() print('iter {}, epoch {}, generator loss {:.3f}'.format(iter, epoch, loss.item())) iter += 1 self.evaluator.evaluate_generator(self.generator) torch.save(self.generator.state_dict(), self.generator_checkpoint_path) def _pretrain_discriminator(self): iter = 0 for epoch in range(self.args.pretrain_discriminator_epochs): self.discriminator.train() for data in self.discriminator_loader: loss = self._train_discriminator(data) print('iter {}, epoch {}, discriminator loss {:.3f}'.format(iter, epoch, loss)) iter += 1 self.evaluator.evaluate_discriminator(generator=self.generator, discriminator=self.discriminator) torch.save(self.discriminator.state_dict(), self.discriminator_checkpoint_path) def _train_gan(self): generator_iter = iter(self.generator_loader) discriminator_iter = iter(self.discriminator_loader) for i in range(self.args.train_gan_iters): print('iter {}'.format(i)) for j in range(1): try: data = next(generator_iter) except StopIteration: generator_iter = iter(self.generator_loader) data = next(generator_iter) result = self._train_generator(data) print('generator loss {:.3f}, fake prob {:.3f}, cider score {:.3f}'.format(result['loss'], result['fake_prob'], result['cider_score'])) for j in range(1): try: data = next(discriminator_iter) except StopIteration: discriminator_iter = iter(self.discriminator_loader) data = next(discriminator_iter) loss = self._train_discriminator(data) print('discriminator loss {:.3f}'.format(loss)) if i != 0 and i % 10000 == 0: self.evaluator.evaluate_generator(self.generator) torch.save(self.generator.state_dict(), self.generator_checkpoint_path) self.evaluator.evaluate_discriminator(generator=self.generator, discriminator=self.discriminator) torch.save(self.discriminator.state_dict(), self.discriminator_checkpoint_path) def _train_generator(self, data): self.generator.train() for name, item in data.items(): data[name] = item.to(self.device) self.generator.zero_grad() probs = self.generator(data['fc_feats'], data['att_feats'], data['att_masks'], data['labels']) loss1 = self.sequence_loss(probs, data['labels']) seqs, probs = self.generator.sample(data['fc_feats'], data['att_feats'], data['att_masks']) greedy_seqs = self.generator.greedy_decode(data['fc_feats'], data['att_feats'], data['att_masks']) reward, fake_prob, score = self._get_reward(data, seqs) baseline, _, _ = self._get_reward(data, greedy_seqs) loss2 = self.reinforce_loss(reward, baseline, probs, seqs) loss = loss1 + loss2 loss.backward() self.generator_optimizer.step() result = { 'loss': loss1.item(), 'fake_prob': fake_prob, 'cider_score': score } return result def _train_discriminator(self, data): self.discriminator.train() for name, item in data.items(): data[name] = item.to(self.device) self.discriminator.zero_grad() real_probs = self.discriminator(data['fc_feats'], data['att_feats'], data['att_masks'], data['labels']) wrong_probs = self.discriminator(data['fc_feats'], data['att_feats'], data['att_masks'], data['wrong_labels']) # generate fake data with torch.no_grad(): fake_seqs, _ = self.generator.sample(data['fc_feats'], data['att_feats'], data['att_masks']) fake_probs = self.discriminator(data['fc_feats'], data['att_feats'], data['att_masks'], fake_seqs) loss = -(0.5 * torch.log(real_probs + 1e-10) + 0.25 * torch.log(1 - wrong_probs + 1e-10) + 0.25 * torch.log(1 - fake_probs + 1e-10)).mean() loss.backward() self.discriminator_optimizer.step() return loss.item() def _get_reward(self, data, seqs): probs = self.discriminator(data['fc_feats'], data['att_feats'], data['att_masks'], seqs) scores = self.cider.get_scores(seqs.cpu().numpy(), data['images'].cpu().numpy()) reward = probs + torch.tensor(scores, dtype=torch.float, device=self.device) fake_prob = probs.mean().item() score = scores.mean() return reward, fake_prob, score
def run(): print('loop') # torch.backends.cudnn.enabled = False device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") # device = torch.device("cpu") # Assuming that we are on a CUDA machine, this should print a CUDA device: print(device) Dx = Discriminator().to(device) Gx = UNet(3, 3).to(device) Dy = Discriminator().to(device) Gy = UNet(3, 3).to(device) ld = False if ld: try: Gx.load_state_dict(torch.load('./genx')) Dx.load_state_dict(torch.load('./fcnx')) Gy.load_state_dict(torch.load('./geny')) Dy.load_state_dict(torch.load('./fcny')) print('net loaded') except Exception as e: print(e) dataset = 'ukiyoe2photo' # A 562 image_path_A = './datasets/' + dataset + '/trainA/*.jpg' image_path_B = './datasets/' + dataset + '/trainB/*.jpg' plt.ion() train_image_paths_A = glob.glob(image_path_A) train_image_paths_B = glob.glob(image_path_B) print(len(train_image_paths_A), len(train_image_paths_B)) b_size = 8 train_dataset_A = CustomDataset(train_image_paths_A, train=True) train_loader_A = torch.utils.data.DataLoader(train_dataset_A, batch_size=b_size, shuffle=True, num_workers=4, pin_memory=False, drop_last=True) train_dataset_B = CustomDataset(train_image_paths_B, True, 562, train=True) train_loader_B = torch.utils.data.DataLoader(train_dataset_B, batch_size=b_size, shuffle=True, num_workers=4, pin_memory=False, drop_last=True) Gx.train() Dx.train() Gy.train() Dy.train() criterion = nn.BCEWithLogitsLoss().to(device) # criterion2 = nn.SmoothL1Loss().to(device) criterion2 = nn.L1Loss().to(device) g_lr = 2e-4 d_lr = 2e-4 optimizer_x = optim.Adam(Gx.parameters(), lr=g_lr, betas=(0.5, 0.999)) optimizer_x_d = optim.Adam(Dx.parameters(), lr=d_lr, betas=(0.5, 0.999)) optimizer_y = optim.Adam(Gy.parameters(), lr=g_lr, betas=(0.5, 0.999)) optimizer_y_d = optim.Adam(Dy.parameters(), lr=d_lr, betas=(0.5, 0.999)) # cp = cropper().to(device) _zero = torch.from_numpy(np.zeros((b_size, 1))).float().to(device) _zero.requires_grad = False _one = torch.from_numpy(np.ones((b_size, 1))).float().to(device) _one.requires_grad = False for epoch in trange(100, desc='epoch'): # loop = tqdm(zip(train_loader_A, train_loader_B), desc='iteration') loop = zip(tqdm(train_loader_A, desc='iteration'), train_loader_B) batch_idx = 0 for data_A, data_B in loop: batch_idx += 1 zero = _zero one = _one _data_A = data_A.to(device) _data_B = data_B.to(device) # Dy loss (A -> B) gen = Gy(_data_A) optimizer_y_d.zero_grad() output2_p = Dy(_data_B.detach()) output_p = Dy(gen.detach()) errD = ( criterion(output2_p - torch.mean(output_p), one.detach()) + criterion(output_p - torch.mean(output2_p), zero.detach())) / 2 errD.backward() optimizer_y_d.step() # Dx loss (B -> A) gen = Gx(_data_B) optimizer_x_d.zero_grad() output2_p = Dx(_data_A.detach()) output_p = Dx(gen.detach()) errD = ( criterion(output2_p - torch.mean(output_p), one.detach()) + criterion(output_p - torch.mean(output2_p), zero.detach())) / 2 errD.backward() optimizer_x_d.step() # Gy loss (A -> B) optimizer_y.zero_grad() gen = Gy(_data_A) output_p = Dy(gen) output2_p = Dy(_data_B.detach()) g_loss = ( criterion(output2_p - torch.mean(output_p), zero.detach()) + criterion(output_p - torch.mean(output2_p), one.detach())) / 2 # Gy cycle loss (B -> A -> B) fA = Gx(_data_B) gen = Gy(fA.detach()) c_loss = criterion2(gen, _data_B) errG = g_loss + c_loss errG.backward() optimizer_y.step() if batch_idx % 10 == 0: fig = plt.figure(1) fig.clf() plt.imshow((np.transpose(_data_B.detach().cpu().numpy()[0], (1, 2, 0)) + 1) / 2) fig.canvas.draw() fig.canvas.flush_events() fig = plt.figure(2) fig.clf() plt.imshow((np.transpose(fA.detach().cpu().numpy()[0], (1, 2, 0)) + 1) / 2) fig.canvas.draw() fig.canvas.flush_events() fig = plt.figure(3) fig.clf() plt.imshow((np.transpose(gen.detach().cpu().numpy()[0], (1, 2, 0)) + 1) / 2) fig.canvas.draw() fig.canvas.flush_events() # Gx loss (B -> A) optimizer_x.zero_grad() gen = Gx(_data_B) output_p = Dx(gen) output2_p = Dx(_data_A.detach()) g_loss = ( criterion(output2_p - torch.mean(output_p), zero.detach()) + criterion(output_p - torch.mean(output2_p), one.detach())) / 2 # Gx cycle loss (A -> B -> A) fB = Gy(_data_A) gen = Gx(fB.detach()) c_loss = criterion2(gen, _data_A) errG = g_loss + c_loss errG.backward() optimizer_x.step() torch.save(Gx.state_dict(), './genx') torch.save(Dx.state_dict(), './fcnx') torch.save(Gy.state_dict(), './geny') torch.save(Dy.state_dict(), './fcny') print('\nFinished Training')
class SAGAN(): def __init__(self, dataloader, configs): # Data Loader self.dataloader = dataloader # model settings & hyperparams self.total_steps = configs.total_steps self.d_iters = configs.d_iters self.g_iters = configs.g_iters self.batch_size = configs.batch_size self.imsize = configs.imsize self.nz = configs.nz self.ngf = configs.ngf self.ndf = configs.ndf self.g_lr = configs.g_lr self.d_lr = configs.d_lr self.beta1 = configs.beta1 self.beta2 = configs.beta2 # instance noise self.inst_noise_sigma = configs.inst_noise_sigma self.inst_noise_sigma_iters = configs.inst_noise_sigma_iters # model logging and saving self.log_step = configs.log_step self.save_epoch = configs.save_epoch self.model_path = configs.model_path self.sample_path = configs.sample_path # pretrained self.pretrained_model = configs.pretrained_model # building self.build_model() # archive of all losses self.ave_d_losses = [] self.ave_d_losses_real = [] self.ave_d_losses_fake = [] self.ave_d_gamma1 = [] self.ave_d_gamma2 = [] self.ave_g_losses = [] self.ave_g_gamma1 = [] self.ave_g_gamma2 = [] if self.pretrained_model: self.load_pretrained() def build_model(self): # initialize Generator and Discriminator self.G = Generator(self.imsize, self.nz, self.ngf).cuda() self.D = Discriminator(self.ndf).cuda() # optimizers self.g_optimizer = optim.Adam(filter( lambda p: p.requires_grad, self.G.parameters()), self.g_lr, [self.beta1, self.beta2]) self.d_optimizer = optim.Adam(filter( lambda p: p.requires_grad, self.D.parameters()), self.d_lr, [self.beta1, self.beta2]) # tensorboard writer self.tb = SummaryWriter() print("Generator Parameters: ", parameters(self.G)) print(self.G) print("Discriminator Parameters: ", parameters(self.D)) print(self.D) def load_pretrained(self): """Loading pretrained model""" checkpoint = torch.load( os.path.join(self.model_path, "{}_sagan.pth".format( self.pretrained_model))) # load models self.G.load_state_dict(checkpoint["gen_state_dict"]) self.D.load_state_dict(checkpoint["disc_state_dict"]) # load optimizers self.g_optimizer.load_state_dict(checkpoint["gen_optimizer"]) self.d_optimizer.load_state_dict(checkpoint["disc_optimizer"]) # load losses self.ave_d_losses = checkpoint["ave_d_losses"] self.ave_d_losses_real = checkpoint["ave_d_losses_real"] self.ave_d_losses_fake = checkpoint["ave_d_losses_fake"] self.ave_d_gamma1 = checkpoint["ave_d_gamma1"] self.ave_d_gamma2 = checkpoint["ave_d_gamma2"] self.ave_g_losses = checkpoint["ave_g_losses"] self.ave_g_gamma1 = checkpoint["ave_g_gamma1"] self.ave_g_gamma2 = checkpoint["ave_g_gamma2"] print("Loading pretrained models (epoch: {})..!".format( self.pretrained_model)) def reset_grad(self): """Reset gradients""" self.g_optimizer.zero_grad() self.d_optimizer.zero_grad() def train(self): step_per_epoch = len(self.dataloader) epochs = int(self.total_steps / step_per_epoch) # fixed z for sampling generator images fixed_z = tensor2var(torch.randn(self.batch_size, self.nz)) print("Initiating Training") print("Epochs: {}, Total Steps: {}, Steps/Epoch: {}". format(epochs, self.total_steps, step_per_epoch)) if self.pretrained_model: start_epoch = self.pretrained_model else: start_epoch = 0 # train layers self.D.train() self.G.train() # total time start_time = time.time() for epoch in range(start_epoch, epochs): # local losses d_losses = [] d_losses_real = [] d_losses_fake = [] d_gamma1 = [] d_gamma2 = [] g_losses = [] g_gamma1 = [] g_gamma2 = [] data_iter = iter(self.dataloader) for step in range(step_per_epoch): # get real images real_images, _ = next(data_iter) real_images = tensor2var(real_images) # Instance noise - make random noise mean (0) and std for injecting inst_noise_mean = torch.full( (real_images.size(0), 3, self.imsize, self.imsize), 0).cuda() inst_noise_std = torch.full( (real_images.size(0), 3, self.imsize, self.imsize), self.inst_noise_sigma).cuda() # Instance noise std is linearly annealed from self.inst_noise_sigma to 0 thru self.inst_noise_sigma_iters inst_noise_sigma_curr = 0 if step > self.inst_noise_sigma_iters else ( 1 - step/self.inst_noise_sigma_iters)*self.inst_noise_sigma inst_noise_std.fill_(inst_noise_sigma_curr) # ================== TRAIN DISCRIMINATOR ================== # for _ in range(self.d_iters): self.reset_grad() # TRAIN REAL # creating instance noise inst_noise = torch.normal( mean=inst_noise_mean, std=inst_noise_std).cuda() # get D output for real images + noise d_real = self.D(real_images + inst_noise) # compute hinge loss of D with real images d_loss_real = loss_hinge_dis_real(d_real) d_loss_real.backward() # TRAIN FAKE # generate fake images and get D output for fake images z = tensor2var(torch.randn(real_images.size(0), self.nz)) fake_images = self.G(z) # creating instance noise inst_noise = torch.normal( mean=inst_noise_mean, std=inst_noise_std).cuda() # adding noise to fake images # get D output for fake images d_fake = self.D(fake_images + inst_noise) # compute hinge loss of D with fake images d_loss_fake = loss_hinge_dis_fake(d_fake) d_loss_fake.backward() d_loss = d_loss_real + d_loss_fake # optimize D self.d_optimizer.step() # ================== TRAIN GENERATOR ================== # for _ in range(self.g_iters): self.reset_grad() # create new latent vector z = tensor2var(torch.randn(real_images.size(0), self.nz)) inst_noise = torch.normal( mean=inst_noise_mean, std=inst_noise_std).cuda() # generate fake images fake_images = self.G(z) g_fake = self.D(fake_images + inst_noise) # compute hinge loss for G g_loss = loss_hinge_gen(g_fake) g_loss.backward() self.g_optimizer.step() # logging step progression if (step+1) % self.log_step == 0: # logging losses and attention d_losses.append(d_loss.item()) d_losses_real.append(d_loss_real.item()) d_losses_fake.append(d_loss_fake.item()) d_gamma1.append(self.D.attn1.gamma.data.item()) d_gamma2.append(self.D.attn2.gamma.data.item()) g_losses.append(g_loss.item()) g_gamma1.append(self.G.attn1.gamma.data.item()) g_gamma2.append(self.G.attn2.gamma.data.item()) # print out elapsed = time.time() - start_time elapsed = str(datetime.timedelta(seconds=elapsed)) print("Elapsed [{}], Epoch: [{}/{}], Step [{}/{}], g_loss: {:.4f}, d_loss: {:.4f}," " d_loss_real: {:.4f}, d_loss_fake: {:.4f}". format(elapsed, epoch+1, epochs, (step + 1), step_per_epoch, g_loss, d_loss, d_loss_real, d_loss_fake)) # logging average losses over epoch self.ave_d_losses.append(mean(d_losses)) self.ave_d_losses_real.append(mean(d_losses_real)) self.ave_d_losses_fake.append(mean(d_losses_fake)) self.ave_d_gamma1.append(mean(d_gamma1)) self.ave_d_gamma2.append(mean(d_gamma2)) self.ave_g_losses.append(mean(g_losses)) self.ave_g_gamma1.append(mean(g_gamma1)) self.ave_g_gamma2.append(mean(g_gamma2)) # adding tensorboard logs self.tb.add_scalar("d loss", self.ave_d_losses[epoch], epoch) self.tb.add_scalar('d real', self.ave_d_losses_real[epoch], epoch) self.tb.add_scalar('d fake', self.ave_d_losses_fake[epoch], epoch) self.tb.add_scalar("g loss", self.ave_g_losses[epoch], epoch) self.tb.add_scalar("g gamma 1", self.ave_g_gamma1[epoch], epoch) self.tb.add_scalar("g gamma 2", self.ave_g_gamma2[epoch], epoch) self.tb.add_scalar("d gamma 1", self.ave_d_gamma1[epoch], epoch) self.tb.add_scalar("d gamma 2", self.ave_d_gamma2[epoch], epoch) # epoch update print("Elapsed [{}], Epoch: [{}/{}], ave_g_loss: {:.4f}, ave_d_loss: {:.4f}," " ave_d_loss_real: {:.4f}, ave_d_loss_fake: {:.4f}," " ave_g_gamma1: {:.4f}, ave_g_gamma2: {:.4f}, ave_d_gamma1: {:.4f}, ave_d_gamma2: {:.4f}". format(elapsed, epoch+1, epochs, self.ave_g_losses[epoch], self.ave_d_losses[epoch], self.ave_d_losses_real[epoch], self.ave_d_losses_fake[epoch], self.ave_g_gamma1[epoch], self.ave_g_gamma2[epoch], self.ave_d_gamma1[epoch], self.ave_d_gamma2[epoch])) # sample images every epoch fake_images = self.G(fixed_z) fake_images = denorm(fake_images.data) save_image(fake_images, os.path.join(self.sample_path, "Epoch {}.png".format(epoch+1))) # save model if (epoch+1) % self.save_epoch == 0: torch.save({ "gen_state_dict": self.G.state_dict(), "disc_state_dict": self.D.state_dict(), "gen_optimizer": self.g_optimizer.state_dict(), "disc_optimizer": self.d_optimizer.state_dict(), "ave_d_losses": self.ave_d_losses, "ave_d_losses_real": self.ave_d_losses_real, "ave_d_losses_fake": self.ave_d_losses_fake, "ave_d_gamma1": self.ave_d_gamma1, "ave_d_gamma2": self.ave_d_gamma2, "ave_g_losses": self.ave_g_losses, "ave_g_gamma1": self.ave_g_gamma1, "ave_g_gamma2": self.ave_g_gamma2 }, os.path.join(self.model_path, "{}_sagan.pth".format(epoch+1))) print("Saving models (epoch {})..!".format(epoch+1)) def plot(self): plt.plot(self.ave_d_losses) plt.plot(self.ave_d_losses_real) plt.plot(self.ave_d_losses_fake) plt.plot(self.ave_g_losses) plt.legend(["d loss", "d real", "d fake", "g loss"], loc="upper left") plt.show() def sample(self, samples): z = tensor2var(torch.randn(samples, self.nz)) images = self.G(z) images = denorm(images.data) # https://pytorch.org/docs/stable/_modules/torchvision/utils.html#save_image grid = make_grid(images, nrow=8, padding=2, pad_value=0, normalize=False, range=None, scale_each=False) # Add 0.5 after unnormalizing to [0, 255] to round to nearest integer ndarr = grid.mul(255).add_(0.5).clamp_(0, 255).permute( 1, 2, 0).to('cpu', torch.uint8).numpy() im = Image.fromarray(ndarr) plt.imshow(im) plt.show()
def _main(): print_gpu_details() device = torch.device("cuda:0" if (torch.cuda.is_available()) else "cpu") train_root = args.train_path image_size = 256 cropped_image_size = 256 print("set image folder") train_set = dset.ImageFolder(root=train_root, transform=transforms.Compose([ transforms.Resize(image_size), transforms.CenterCrop(cropped_image_size), transforms.ToTensor() ])) normalizer_clf = transforms.Compose([ transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ]) normalizer_discriminator = transforms.Compose([ transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5]) ]) print('set data loader') train_loader = torch.utils.data.DataLoader(train_set, batch_size=args.batch_size, shuffle=True, num_workers=4, pin_memory=True, drop_last=True) # Network creation classifier = torch.load(args.classifier_path) classifier.eval() generator = Generator(gen_type=args.gen_type) discriminator = Discriminator(args.discriminator_norm, dis_type=args.gen_type) # init weights if args.generator_path is not None: generator.load_state_dict(torch.load(args.generator_path)) else: generator.init_weights() if args.discriminator_path is not None: discriminator.load_state_dict(torch.load(args.discriminator_path)) else: discriminator.init_weights() classifier.to(device) generator.to(device) discriminator.to(device) # losses + optimizers criterion_discriminator, criterion_generator = get_wgan_losses_fn() criterion_features = nn.L1Loss() criterion_diversity_n = nn.L1Loss() criterion_diversity_d = nn.L1Loss() generator_optimizer = optim.Adam(generator.parameters(), lr=args.lr, betas=(0.5, 0.999)) discriminator_optimizer = optim.Adam(discriminator.parameters(), lr=args.lr, betas=(0.5, 0.999)) num_of_epochs = args.epochs starting_time = time.time() iterations = 0 # creating dirs for keeping models checkpoint, temp created images, and loss summary outputs_dir = os.path.join('wgan-gp_models', args.model_name) if not os.path.isdir(outputs_dir): os.makedirs(outputs_dir, exist_ok=True) temp_results_dir = os.path.join(outputs_dir, 'temp_results') if not os.path.isdir(temp_results_dir): os.mkdir(temp_results_dir) models_dir = os.path.join(outputs_dir, 'models_checkpoint') if not os.path.isdir(models_dir): os.mkdir(models_dir) writer = tensorboardX.SummaryWriter(os.path.join(outputs_dir, 'summaries')) z = torch.randn(args.batch_size, 128, 1, 1).to(device) # a fixed noise for sampling z2 = torch.randn(args.batch_size, 128, 1, 1).to(device) # a fixed noise for diversity sampling fixed_features = 0 fixed_masks = 0 fixed_features_diversity = 0 first_iter = True print("Starting Training Loop...") for epoch in range(num_of_epochs): for data in train_loader: train_type = random.choices([1, 2], [args.train1_prob, 1-args.train1_prob]) # choose train type iterations += 1 if iterations % 30 == 1: print('epoch:', epoch, ', iter', iterations, 'start, time =', time.time() - starting_time, 'seconds') starting_time = time.time() images, _ = data images = images.to(device) # change to gpu tensor images_discriminator = normalizer_discriminator(images) images_clf = normalizer_clf(images) _, features = classifier(images_clf) if first_iter: # save batch of images to keep track of the model process first_iter = False fixed_features = [torch.clone(features[x]) for x in range(len(features))] fixed_masks = [torch.ones(features[x].shape, device=device) for x in range(len(features))] fixed_features_diversity = [torch.clone(features[x]) for x in range(len(features))] for i in range(len(features)): for j in range(fixed_features_diversity[i].shape[0]): fixed_features_diversity[i][j] = fixed_features_diversity[i][j % 8] grid = vutils.make_grid(images_discriminator, padding=2, normalize=True, nrow=8) vutils.save_image(grid, os.path.join(temp_results_dir, 'original_images.jpg')) orig_images_diversity = torch.clone(images_discriminator) for i in range(orig_images_diversity.shape[0]): orig_images_diversity[i] = orig_images_diversity[i % 8] grid = vutils.make_grid(orig_images_diversity, padding=2, normalize=True, nrow=8) vutils.save_image(grid, os.path.join(temp_results_dir, 'original_images_diversity.jpg')) # Select a features layer to train on features_to_train = random.randint(1, len(features) - 2) if args.fixed_layer is None else args.fixed_layer # Set masks masks = [features[i].clone() for i in range(len(features))] setMasksPart1(masks, device, features_to_train) if train_type == 1 else setMasksPart2(masks, device, features_to_train) discriminator_loss_dict = train_discriminator(generator, discriminator, criterion_discriminator, discriminator_optimizer, images_discriminator, features, masks) for k, v in discriminator_loss_dict.items(): writer.add_scalar('D/%s' % k, v.data.cpu().numpy(), global_step=iterations) if iterations % 30 == 1: print('{}: {:.6f}'.format(k, v)) if iterations % args.discriminator_steps == 1: generator_loss_dict = train_generator(generator, discriminator, criterion_generator, generator_optimizer, images.shape[0], features, criterion_features, features_to_train, classifier, normalizer_clf, criterion_diversity_n, criterion_diversity_d, masks, train_type) for k, v in generator_loss_dict.items(): writer.add_scalar('G/%s' % k, v.data.cpu().numpy(), global_step=iterations//5 + 1) if iterations % 30 == 1: print('{}: {:.6f}'.format(k, v)) # Save generator and discriminator weights every 1000 iterations if iterations % 1000 == 1: torch.save(generator.state_dict(), models_dir + '/' + args.model_name + 'G') torch.save(discriminator.state_dict(), models_dir + '/' + args.model_name + 'D') # Save temp results if args.keep_temp_results: if iterations < 10000 and iterations % 1000 == 1 or iterations % 2000 == 1: # regular sampling (batch of different images) first_features = True fake_images = None fake_images_diversity = None for i in range(1, 5): one_layer_mask = isolate_layer(fixed_masks, i, device) if first_features: first_features = False fake_images = sample(generator, z, fixed_features, one_layer_mask) fake_images_diversity = sample(generator, z, fixed_features_diversity, one_layer_mask) else: tmp_fake_images = sample(generator, z, fixed_features, one_layer_mask) fake_images = torch.vstack((fake_images, tmp_fake_images)) tmp_fake_images = sample(generator, z2, fixed_features_diversity, one_layer_mask) fake_images_diversity = torch.vstack((fake_images_diversity, tmp_fake_images)) grid = vutils.make_grid(fake_images, padding=2, normalize=True, nrow=8) vutils.save_image(grid, os.path.join(temp_results_dir, 'res_iter_{}.jpg'.format(iterations // 1000))) # diversity sampling (8 different images each with few different noises) grid = vutils.make_grid(fake_images_diversity, padding=2, normalize=True, nrow=8) vutils.save_image(grid, os.path.join(temp_results_dir, 'div_iter_{}.jpg'.format(iterations // 1000))) if iterations % 20000 == 1: torch.save(generator.state_dict(), models_dir + '/' + args.model_name + 'G_' + str(iterations // 15000)) torch.save(discriminator.state_dict(), models_dir + '/' + args.model_name + 'D_' + str(iterations // 15000))
def train(config): genAB = UNet(3, 3, bilinear=config.model.bilinear_upsample).cuda() init_weights(genAB, 'normal') genBA = UNet(3, 3, bilinear=config.model.bilinear_upsample).cuda() init_weights(genBA, 'normal') discrA = Discriminator(3).cuda() init_weights(discrA, 'normal') discrB = Discriminator(3).cuda() init_weights(discrB, 'normal') writer = SummaryWriter(config.name) data_train, data_test = datasets_by_name(config.dataset.name, config.dataset) train_dataloader = DataLoader(data_train, batch_size=config.bs, shuffle=True, num_workers=config.num_workers) test_dataloader = DataLoader(data_test, batch_size=config.bs, shuffle=True, num_workers=config.num_workers) idt_loss = nn.L1Loss() cycle_consistency = nn.L1Loss() l2_loss = nn.MSELoss() discriminator_loss = nn.BCELoss() lambda_idt, lambda_C, lambda_D = config.loss.lambda_idt, config.loss.lambda_C, config.loss.lambda_D optG = torch.optim.Adam(itertools.chain(genAB.parameters(), genBA.parameters()), lr=config.train.lr, betas=(config.train.beta1, 0.999)) optD = torch.optim.Adam(itertools.chain(discrA.parameters(), discrB.parameters()), lr=config.train.lr, betas=(config.train.beta1, 0.999)) genAB, genBA, discrA, discrB, optG, optD, start_epoch = load_if_exsists( config, genAB, genBA, discrA, discrB, optG, optD) for epoch in range(start_epoch, config.train.epochs): set_train([genAB, genBA, discrA, discrB]) set_requires_grad([genAB, genBA, discrA, discrB], True) for i, (batch_A, batch_B) in enumerate(tqdm(train_dataloader)): batch_A, batch_B = batch_A.cuda(), batch_B.cuda() optG.zero_grad() loss_G, loss_D = 0, 0 fake_B = genAB(batch_A) cycle_A = genBA(fake_B) fake_A = genBA(batch_B) cycle_B = genAB(fake_A) if lambda_idt > 0: loss_G += idt_loss(fake_B, batch_B) * lambda_idt loss_G += idt_loss(fake_A, batch_A) * lambda_idt if lambda_C > 0: loss_G += cycle_consistency(cycle_A, batch_A) * lambda_C loss_G += cycle_consistency(cycle_B, batch_B) * lambda_C if lambda_D > 0: set_requires_grad([discrA, discrB], False) discr_feedbackA = discrA(fake_A) discr_feedbackB = discrB(fake_B) loss_G += discriminator_loss( discr_feedbackA, torch.ones_like(discr_feedbackA)) * lambda_D loss_G += discriminator_loss( discr_feedbackB, torch.ones_like(discr_feedbackB)) * lambda_D loss_G.backward() torch.nn.utils.clip_grad_norm_( itertools.chain(genAB.parameters(), genBA.parameters()), 15) optG.step() if lambda_D > 0: set_requires_grad([discrA, discrB], True) loss_D_fake, loss_D_true = 0, 0 optD.zero_grad() logits = discrA(fake_A.detach()) loss_D_fake += discriminator_loss(logits, torch.zeros_like(logits)) logits = discrB(fake_B.detach()) loss_D_fake += discriminator_loss(logits, torch.zeros_like(logits)) loss_D_fake.backward() torch.nn.utils.clip_grad_norm_( itertools.chain(discrA.parameters(), discrB.parameters()), 15) optD.step() optD.zero_grad() logits = discrA(batch_A) loss_D_true += discriminator_loss(logits, torch.ones_like(logits)) logits = discrB(batch_B) loss_D_true += discriminator_loss(logits, torch.ones_like(logits)) loss_D_true.backward() torch.nn.utils.clip_grad_norm_( itertools.chain(discrA.parameters(), discrB.parameters()), 15) optD.step() loss_D = loss_D_fake + loss_D_true if (i % config.train.verbose_period == 0): writer.add_scalar('train/loss_G', loss_G.item(), len(train_dataloader) * epoch + i) writer.add_scalar('train/pixel_error_A', l2_loss(fake_A, batch_A).mean().item(), len(train_dataloader) * epoch + i) writer.add_scalar('train/pixel_error_B', l2_loss(fake_B, batch_B).mean().item(), len(train_dataloader) * epoch + i) if lambda_D > 0: writer.add_scalar('train/loss_D', loss_D.item(), len(train_dataloader) * epoch + i) writer.add_scalar('train/mean_D_A', discr_feedbackA.mean().item(), len(train_dataloader) * epoch + i) writer.add_scalar('train/mean_D_B', discr_feedbackB.mean().item(), len(train_dataloader) * epoch + i) for batch_i in range(fake_A.shape[0]): concat = (torch.cat([fake_A[batch_i], batch_B[batch_i]], dim=-1) + 1.) / 2. writer.add_image('train/fake_A_' + str(batch_i), concat, len(train_dataloader) * epoch + i) for batch_i in range(fake_B.shape[0]): concat = (torch.cat([fake_B[batch_i], batch_A[batch_i]], dim=-1) + 1.) / 2. writer.add_image('train/fake_B_' + str(batch_i), concat, len(train_dataloader) * epoch + i) if not config.validate: continue set_eval([genAB, genBA, discrA, discrB]) set_requires_grad([genAB, genBA, discrA, discrB], False) loss_G, loss_D, discr_feedbackA_mean, discr_feedbackB_mean = 0, 0, 0, 0 pixel_error_A, pixel_error_B = 0, 0 for i, (batch_A, batch_B) in enumerate(tqdm(test_dataloader)): batch_A, batch_B = batch_A.cuda(), batch_B.cuda() fake_B = genAB(batch_A) cycle_A = genBA(fake_B) fake_A = genBA(batch_B) cycle_B = genAB(fake_A) pixel_error_A += l2_loss(fake_A, batch_A).mean() pixel_error_B += l2_loss(fake_B, batch_B).mean() if lambda_idt > 0: loss_G += idt_loss(fake_B, batch_B) * lambda_idt loss_G += idt_loss(fake_A, batch_A) * lambda_idt if lambda_C > 0: loss_G += cycle_consistency(cycle_A, batch_A) * lambda_C loss_G += cycle_consistency(cycle_B, batch_B) * lambda_C if lambda_D > 0: discr_feedbackA = discrA(fake_A) discr_feedbackB = discrB(fake_B) loss_G += discriminator_loss( discr_feedbackA, torch.ones_like(discr_feedbackA)) * lambda_D loss_G += discriminator_loss( discr_feedbackB, torch.ones_like(discr_feedbackB)) * lambda_D discr_feedbackA_mean += discr_feedbackA.mean() discr_feedbackB_mean += discr_feedbackB.mean() if lambda_D > 0: loss_D_fake, loss_D_true = 0, 0 logits = discrA(fake_A.detach()) loss_D_fake += discriminator_loss(logits, torch.zeros_like(logits)) logits = discrB(fake_B.detach()) loss_D_fake += discriminator_loss(logits, torch.zeros_like(logits)) logits = discrA(batch_A) loss_D_true += discriminator_loss(logits, torch.ones_like(logits)) logits = discrB(batch_B) loss_D_true += discriminator_loss(logits, torch.ones_like(logits)) loss_D += loss_D_fake + loss_D_true if i == 0: for batch_i in range(fake_A.shape[0]): concat = (torch.cat([fake_A[batch_i], batch_B[batch_i]], dim=-1) + 1.) / 2. writer.add_image('val/fake_A_' + str(batch_i), concat, epoch) for batch_i in range(fake_B.shape[0]): concat = (torch.cat([fake_B[batch_i], batch_A[batch_i]], dim=-1) + 1.) / 2. writer.add_image('val/fake_B_' + str(batch_i), concat, epoch) loss_G /= len(test_dataloader) pixel_error_A /= len(test_dataloader) pixel_error_B /= len(test_dataloader) writer.add_scalar('val/loss_G', loss_G.item(), epoch) writer.add_scalar('val/pixel_error_A', pixel_error_A.item(), epoch) writer.add_scalar('val/pixel_error_B', pixel_error_B.item(), epoch) if lambda_D > 0: loss_D /= len(test_dataloader) discr_feedbackA_mean /= len(test_dataloader) discr_feedbackB_mean /= len(test_dataloader) writer.add_scalar('val/loss_D', loss_D.item(), epoch) writer.add_scalar('val/mean_D_A', discr_feedbackA_mean.item(), epoch) writer.add_scalar('val/mean_D_B', discr_feedbackB_mean.item(), epoch) torch.save( { 'genAB': genAB.state_dict(), 'genBA': genBA.state_dict(), 'discrA': discrA.state_dict(), 'discrB': discrB.state_dict(), 'optG': optG.state_dict(), 'optD': optD.state_dict(), 'epoch': epoch }, os.path.join(config.name, 'model.pth'))
if opt.cuda: rewards = rewards.cuda() rewards = torch.exp(rewards).contiguous().view((-1, )) prob = generators[i].forward(inputs) loss = gen_gan_losses[i](prob, targets, rewards) gen_gan_optm[i].zero_grad() loss.backward() gen_gan_optm[i].step() if total_batch % 10 == 0 or total_batch == TOTAL_BATCH - 1: for generator in generators: print('Saving generator {} with bleu_4: {}'.format( generator.name, bleu_4(TEXT, corpus, generator, g_sequence_len, count=100))) torch.save( generator.state_dict(), CHECKPOINT_PATH + 'generator_seqgan_{}.gen'.format(generator.name)) torch.save(discriminator.state_dict(), CHECKPOINT_PATH + 'discriminator_seqgan.dis') for rollout in rollouts: rollout.update_params() for _ in range(1): loss, acc = train_discriminator(discriminator, generators, real_data_iterator, dis_criterion, dis_optimizer) print('Epoch [%d], loss: %f, accuracy: %f' % (total_batch, loss, acc)) # if __name__ == '__main__': # main()
def main(): # # -------------------- Data -------------------- num_workers = 8 # number of subprocesses to use for data loading batch_size = 64 # how many samples per batch to load transform = transforms.ToTensor() # convert data to torch.FloatTensor train_data = datasets.MNIST(root='../data', train=True, download=True, transform=transform) train_loader = torch.utils.data.DataLoader(train_data, batch_size=batch_size, num_workers=num_workers) # # Obtain one batch of training images # dataiter = iter(train_loader) # images, labels = dataiter.next() # images = images.numpy() # # Get one image from the batch for visualization # img = np.squeeze(images[0]) # fig = plt.figure(figsize=(3, 3)) # ax = fig.add_subplot(111) # ax.imshow(img, cmap='gray') # plt.show() # # -------------------- Discriminator and Generator -------------------- # Discriminator hyperparams input_size = 784 # Size of input image to discriminator (28*28) d_output_size = 1 # Size of discriminator output (real or fake) d_hidden_size = 32 # Size of last hidden layer in the discriminator # Generator hyperparams z_size = 100 # Size of latent vector to give to generator g_output_size = 784 # Size of discriminator output (generated image) g_hidden_size = 32 # Size of first hidden layer in the generator # Instantiate discriminator and generator D = Discriminator(input_size, d_hidden_size, d_output_size) G = Generator(z_size, g_hidden_size, g_output_size) # # -------------------- Optimizers and Criterion -------------------- # Training hyperparams num_epochs = 100 print_every = 400 lr = 0.002 # Create optimizers for the discriminator and generator, respectively d_optimizer = optim.Adam(D.parameters(), lr) g_optimizer = optim.Adam(G.parameters(), lr) losses = [] # keep track of generated "fake" samples criterion = nn.BCEWithLogitsLoss() # -------------------- Training -------------------- D.train() G.train() # Get some fixed data for sampling. These are images that are held # constant throughout training, and allow us to inspect the model's performance sample_size = 16 fixed_z = np.random.uniform(-1, 1, size=(sample_size, z_size)) fixed_z = torch.from_numpy(fixed_z).float() samples = [] # keep track of loss for epoch in range(num_epochs): for batch_i, (real_images, _) in enumerate(train_loader): batch_size = real_images.size(0) # Important rescaling step real_images = real_images * 2 - 1 # rescale input images from [0,1) to [-1, 1) # Generate fake images, used for both discriminator and generator z = np.random.uniform(-1, 1, size=(batch_size, z_size)) z = torch.from_numpy(z).float() fake_images = G(z) real_labels = torch.ones(batch_size) fake_labels = torch.zeros(batch_size) # ============================================ # TRAIN THE DISCRIMINATOR # ============================================ d_optimizer.zero_grad() # 1. Train with real images # Compute the discriminator losses on real images D_real = D(real_images) d_real_loss = real_loss(criterion, D_real, real_labels, smooth=True) # 2. Train with fake images # Compute the discriminator losses on fake images # ------------------------------------------------------- # ATTENTION: # *.detach(), thus, generator is fixed when we optimize # the discriminator # ------------------------------------------------------- D_fake = D(fake_images.detach()) d_fake_loss = fake_loss(criterion, D_fake, fake_labels) # 3. Add up loss and perform backprop d_loss = (d_real_loss + d_fake_loss) * 0.5 d_loss.backward() d_optimizer.step() # ========================================= # TRAIN THE GENERATOR # ========================================= g_optimizer.zero_grad() # Make the discriminator fixed when optimizing the generator set_model_gradient(D, False) # 1. Train with fake images and flipped labels # Compute the discriminator losses on fake images using flipped labels! G_D_fake = D(fake_images) g_loss = real_loss(criterion, G_D_fake, real_labels) # use real loss to flip labels # 2. Perform backprop g_loss.backward() g_optimizer.step() # Make the discriminator require_grad=True after optimizing the generator set_model_gradient(D, True) # ========================================= # Print some loss stats # ========================================= if batch_i % print_every == 0: print( 'Epoch [{:5d}/{:5d}] | d_loss: {:6.4f} | g_loss: {:6.4f}'. format(epoch + 1, num_epochs, d_loss.item(), g_loss.item())) # AFTER EACH EPOCH losses.append((d_loss.item(), g_loss.item())) # generate and save sample, fake images G.eval() # eval mode for generating samples samples_z = G(fixed_z) samples.append(samples_z) view_samples(-1, samples, "last_sample.png") G.train() # back to train mode # Save models and training generator samples torch.save(G.state_dict(), "G.pth") torch.save(D.state_dict(), "D.pth") with open('train_samples.pkl', 'wb') as f: pkl.dump(samples, f) # Plot the loss curve fig, ax = plt.subplots() losses = np.array(losses) plt.plot(losses.T[0], label='Discriminator') plt.plot(losses.T[1], label='Generator') plt.title("Training Losses") plt.legend() plt.savefig("loss.png") plt.show()
for epoch in range(num_epochs): for batch_i, (real_images, _) in enumerate(train_loader): batch_size = real_images.size(0) real_images = scale(real_images) d_loss = train_discriminator(real_images, d_optim, batch_size, z_size) g_loss = train_generator(g_optim, batch_size, z_size) # Print some loss stats if batch_i % print_every == 0: # print discriminator and generator loss print('Epoch [{:5d}/{:5d}] | d_loss: {:6.4f} | g_loss: {:6.4f}'. format(epoch + 1, num_epochs, d_loss.item(), g_loss.item())) losses.append((d_loss.item(), g_loss.item())) # generate and save sample, fake images G.eval() # eval mode for generating samples samples_z = G(fixed_z) samples.append(samples_z) G.train() # back to train mode torch.save(D.state_dict(), './D.state') torch.save(G.state_dict(), './G.state') with open('train_samples.pkl', 'wb') as f: pkl.dump(samples, f) generate_plot(losses)
def main(args): use_cuda = (len(args.gpuid) >= 1) print("{0} GPU(s) are available".format(cuda.device_count())) print("======printing args========") print(args) print("=================================") # Load dataset splits = ['train', 'valid'] if data.has_binary_files(args.data, splits): print("Loading bin dataset") dataset = data.load_dataset(args.data, splits, args.src_lang, args.trg_lang, args.fixed_max_len) #args.data, splits, args.src_lang, args.trg_lang) else: print(f"Loading raw text dataset {args.data}") dataset = data.load_raw_text_dataset(args.data, splits, args.src_lang, args.trg_lang, args.fixed_max_len) #args.data, splits, args.src_lang, args.trg_lang) if args.src_lang is None or args.trg_lang is None: # record inferred languages in args, so that it's saved in checkpoints args.src_lang, args.trg_lang = dataset.src, dataset.dst print('| [{}] dictionary: {} types'.format(dataset.src, len(dataset.src_dict))) print('| [{}] dictionary: {} types'.format(dataset.dst, len(dataset.dst_dict))) for split in splits: print('| {} {} {} examples'.format(args.data, split, len(dataset.splits[split]))) g_logging_meters = OrderedDict() g_logging_meters['train_loss'] = AverageMeter() g_logging_meters['valid_loss'] = AverageMeter() g_logging_meters['train_acc'] = AverageMeter() g_logging_meters['valid_acc'] = AverageMeter() g_logging_meters['bsz'] = AverageMeter() # sentences per batch d_logging_meters = OrderedDict() d_logging_meters['train_loss'] = AverageMeter() d_logging_meters['valid_loss'] = AverageMeter() d_logging_meters['train_acc'] = AverageMeter() d_logging_meters['valid_acc'] = AverageMeter() d_logging_meters['bsz'] = AverageMeter() # sentences per batch # Set model parameters args.encoder_embed_dim = 1000 args.encoder_layers = 4 args.encoder_dropout_out = 0 args.decoder_embed_dim = 1000 args.decoder_layers = 4 args.decoder_out_embed_dim = 1000 args.decoder_dropout_out = 0 args.bidirectional = False # try to load generator model g_model_path = 'checkpoints/generator/best_gmodel.pt' if not os.path.exists(g_model_path): print("Start training generator!") train_g(args, dataset) assert os.path.exists(g_model_path) generator = LSTMModel(args, dataset.src_dict, dataset.dst_dict, use_cuda=use_cuda) model_dict = generator.state_dict() pretrained_dict = torch.load(g_model_path) #print(f"First dict: {pretrained_dict}") # 1. filter out unnecessary keys pretrained_dict = { k: v for k, v in pretrained_dict.items() if k in model_dict } #print(f"Second dict: {pretrained_dict}") # 2. overwrite entries in the existing state dict model_dict.update(pretrained_dict) #print(f"model dict: {model_dict}") # 3. load the new state dict generator.load_state_dict(model_dict) print("Generator has successfully loaded!") # try to load discriminator model d_model_path = 'checkpoints/discriminator/best_dmodel.pt' if not os.path.exists(d_model_path): print("Start training discriminator!") train_d(args, dataset) assert os.path.exists(d_model_path) discriminator = Discriminator(args, dataset.src_dict, dataset.dst_dict, use_cuda=use_cuda) model_dict = discriminator.state_dict() pretrained_dict = torch.load(d_model_path) # 1. filter out unnecessary keys pretrained_dict = { k: v for k, v in pretrained_dict.items() if k in model_dict } # 2. overwrite entries in the existing state dict model_dict.update(pretrained_dict) # 3. load the new state dict discriminator.load_state_dict(model_dict) print("Discriminator has successfully loaded!") #return print("starting main training loop") torch.autograd.set_detect_anomaly(True) if use_cuda: if torch.cuda.device_count() > 1: discriminator = torch.nn.DataParallel(discriminator).cuda() generator = torch.nn.DataParallel(generator).cuda() else: generator.cuda() discriminator.cuda() else: discriminator.cpu() generator.cpu() # adversarial training checkpoints saving path if not os.path.exists('checkpoints/joint'): os.makedirs('checkpoints/joint') checkpoints_path = 'checkpoints/joint/' # define loss function g_criterion = torch.nn.NLLLoss(size_average=False, ignore_index=dataset.dst_dict.pad(), reduce=True) d_criterion = torch.nn.BCEWithLogitsLoss() pg_criterion = PGLoss(ignore_index=dataset.dst_dict.pad(), size_average=True, reduce=True) # fix discriminator word embedding (as Wu et al. do) for p in discriminator.embed_src_tokens.parameters(): p.requires_grad = False for p in discriminator.embed_trg_tokens.parameters(): p.requires_grad = False # define optimizer g_optimizer = eval("torch.optim." + args.g_optimizer)(filter( lambda x: x.requires_grad, generator.parameters()), args.g_learning_rate) d_optimizer = eval("torch.optim." + args.d_optimizer)( filter(lambda x: x.requires_grad, discriminator.parameters()), args.d_learning_rate, momentum=args.momentum, nesterov=True) # start joint training best_dev_loss = math.inf num_update = 0 # main training loop for epoch_i in range(1, args.epochs + 1): logging.info("At {0}-th epoch.".format(epoch_i)) # seed = args.seed + epoch_i # torch.manual_seed(seed) max_positions_train = (args.fixed_max_len, args.fixed_max_len) # Initialize dataloader, starting at batch_offset itr = dataset.train_dataloader( 'train', max_tokens=args.max_tokens, max_sentences=args.joint_batch_size, max_positions=max_positions_train, # seed=seed, epoch=epoch_i, sample_without_replacement=args.sample_without_replacement, sort_by_source_size=(epoch_i <= args.curriculum), shard_id=args.distributed_rank, num_shards=args.distributed_world_size, ) # reset meters for key, val in g_logging_meters.items(): if val is not None: val.reset() for key, val in d_logging_meters.items(): if val is not None: val.reset() # set training mode generator.train() discriminator.train() update_learning_rate(num_update, 8e4, args.g_learning_rate, args.lr_shrink, g_optimizer) for i, sample in enumerate(itr): if use_cuda: # wrap input tensors in cuda tensors sample = utils.make_variable(sample, cuda=cuda) ## part I: use gradient policy method to train the generator # use policy gradient training when rand > 50% rand = random.random() if rand >= 0.5: # policy gradient training generator.decoder.is_testing = True sys_out_batch, prediction, _ = generator(sample) generator.decoder.is_testing = False with torch.no_grad(): n_i = sample['net_input']['src_tokens'] #print(f"net input:\n{n_i}, pred: \n{prediction}") reward = discriminator( sample['net_input']['src_tokens'], prediction) # dataset.dst_dict.pad()) train_trg_batch = sample['target'] #print(f"sys_out_batch: {sys_out_batch.shape}:\n{sys_out_batch}") pg_loss = pg_criterion(sys_out_batch, train_trg_batch, reward, use_cuda) # logging.debug("G policy gradient loss at batch {0}: {1:.3f}, lr={2}".format(i, pg_loss.item(), g_optimizer.param_groups[0]['lr'])) g_optimizer.zero_grad() pg_loss.backward() torch.nn.utils.clip_grad_norm(generator.parameters(), args.clip_norm) g_optimizer.step() # oracle valid _, _, loss = generator(sample) sample_size = sample['target'].size( 0) if args.sentence_avg else sample['ntokens'] logging_loss = loss.data / sample_size / math.log(2) g_logging_meters['train_loss'].update(logging_loss, sample_size) logging.debug( "G MLE loss at batch {0}: {1:.3f}, lr={2}".format( i, g_logging_meters['train_loss'].avg, g_optimizer.param_groups[0]['lr'])) else: # MLE training #print(f"printing sample: \n{sample}") _, _, loss = generator(sample) sample_size = sample['target'].size( 0) if args.sentence_avg else sample['ntokens'] nsentences = sample['target'].size(0) logging_loss = loss.data / sample_size / math.log(2) g_logging_meters['bsz'].update(nsentences) g_logging_meters['train_loss'].update(logging_loss, sample_size) logging.debug( "G MLE loss at batch {0}: {1:.3f}, lr={2}".format( i, g_logging_meters['train_loss'].avg, g_optimizer.param_groups[0]['lr'])) g_optimizer.zero_grad() loss.backward() # all-reduce grads and rescale by grad_denom for p in generator.parameters(): if p.requires_grad: p.grad.data.div_(sample_size) torch.nn.utils.clip_grad_norm(generator.parameters(), args.clip_norm) g_optimizer.step() num_update += 1 # part II: train the discriminator bsz = sample['target'].size(0) src_sentence = sample['net_input']['src_tokens'] # train with half human-translation and half machine translation true_sentence = sample['target'] true_labels = Variable( torch.ones(sample['target'].size(0)).float()) with torch.no_grad(): generator.decoder.is_testing = True _, prediction, _ = generator(sample) generator.decoder.is_testing = False fake_sentence = prediction fake_labels = Variable( torch.zeros(sample['target'].size(0)).float()) trg_sentence = torch.cat([true_sentence, fake_sentence], dim=0) labels = torch.cat([true_labels, fake_labels], dim=0) indices = np.random.permutation(2 * bsz) trg_sentence = trg_sentence[indices][:bsz] labels = labels[indices][:bsz] if use_cuda: labels = labels.cuda() disc_out = discriminator(src_sentence, trg_sentence) #, dataset.dst_dict.pad()) #print(f"disc out: {disc_out.shape}, labels: {labels.shape}") #print(f"labels: {labels}") d_loss = d_criterion(disc_out, labels.long()) acc = torch.sum(torch.Sigmoid() (disc_out).round() == labels).float() / len(labels) d_logging_meters['train_acc'].update(acc) d_logging_meters['train_loss'].update(d_loss) # logging.debug("D training loss {0:.3f}, acc {1:.3f} at batch {2}: ".format(d_logging_meters['train_loss'].avg, # d_logging_meters['train_acc'].avg, # i)) d_optimizer.zero_grad() d_loss.backward() d_optimizer.step() # validation # set validation mode generator.eval() discriminator.eval() # Initialize dataloader max_positions_valid = (args.fixed_max_len, args.fixed_max_len) itr = dataset.eval_dataloader( 'valid', max_tokens=args.max_tokens, max_sentences=args.joint_batch_size, max_positions=max_positions_valid, skip_invalid_size_inputs_valid_test=True, descending=True, # largest batch first to warm the caching allocator shard_id=args.distributed_rank, num_shards=args.distributed_world_size, ) # reset meters for key, val in g_logging_meters.items(): if val is not None: val.reset() for key, val in d_logging_meters.items(): if val is not None: val.reset() for i, sample in enumerate(itr): with torch.no_grad(): if use_cuda: sample['id'] = sample['id'].cuda() sample['net_input']['src_tokens'] = sample['net_input'][ 'src_tokens'].cuda() sample['net_input']['src_lengths'] = sample['net_input'][ 'src_lengths'].cuda() sample['net_input']['prev_output_tokens'] = sample[ 'net_input']['prev_output_tokens'].cuda() sample['target'] = sample['target'].cuda() # generator validation _, _, loss = generator(sample) sample_size = sample['target'].size( 0) if args.sentence_avg else sample['ntokens'] loss = loss / sample_size / math.log(2) g_logging_meters['valid_loss'].update(loss, sample_size) logging.debug("G dev loss at batch {0}: {1:.3f}".format( i, g_logging_meters['valid_loss'].avg)) # discriminator validation bsz = sample['target'].size(0) src_sentence = sample['net_input']['src_tokens'] # train with half human-translation and half machine translation true_sentence = sample['target'] true_labels = Variable( torch.ones(sample['target'].size(0)).float()) with torch.no_grad(): generator.decoder.is_testing = True _, prediction, _ = generator(sample) generator.decoder.is_testing = False fake_sentence = prediction fake_labels = Variable( torch.zeros(sample['target'].size(0)).float()) trg_sentence = torch.cat([true_sentence, fake_sentence], dim=0) labels = torch.cat([true_labels, fake_labels], dim=0) indices = np.random.permutation(2 * bsz) trg_sentence = trg_sentence[indices][:bsz] labels = labels[indices][:bsz] if use_cuda: labels = labels.cuda() disc_out = discriminator(src_sentence, trg_sentence, dataset.dst_dict.pad()) d_loss = d_criterion(disc_out, labels) acc = torch.sum(torch.Sigmoid()(disc_out).round() == labels).float() / len(labels) d_logging_meters['valid_acc'].update(acc) d_logging_meters['valid_loss'].update(d_loss) # logging.debug("D dev loss {0:.3f}, acc {1:.3f} at batch {2}".format(d_logging_meters['valid_loss'].avg, # d_logging_meters['valid_acc'].avg, i)) torch.save(generator, open( checkpoints_path + "joint_{0:.3f}.epoch_{1}.pt".format( g_logging_meters['valid_loss'].avg, epoch_i), 'wb'), pickle_module=dill) if g_logging_meters['valid_loss'].avg < best_dev_loss: best_dev_loss = g_logging_meters['valid_loss'].avg torch.save(generator, open(checkpoints_path + "best_gmodel.pt", 'wb'), pickle_module=dill)
d_loss.data[0], g_loss.data[0], real_score, fake_score, time.time() - start_time)) record( name='loss_d', value=d_loss.data.cpu().numpy(), data_type='plot', ) record( name='loss_g', value=g_loss.data.cpu().numpy(), data_type='plot', ) log_visdom() counter += 1 d_cost_avg /= num_batch g_cost_avg /= num_batch # Save weights every 3 epoch if (current_epoch + 1) % 3 == 0: print('Epoch:', current_epoch, ' train_loss->', (d_cost_avg, g_cost_avg)) torch.save(generator.state_dict(), './generator.pkl') torch.save(discriminator.state_dict(), './discriminator.pkl') predict(generator, validation_sample, current_epoch, DIR_TO_SAVE) torch.save(generator.state_dict(), './generator.pkl') torch.save(discriminator.state_dict(), './discriminator.pkl') print('Done')
def train(opt): netG_A2B = Unet2(3, 3) netG_B2A = Unet2(3, 3) netD_A = Discriminator(3) netD_B = Discriminator(3) if opt.use_cuda: netG_A2B = netG_A2B.cuda() netG_B2A = netG_B2A.cuda() netD_A = netD_A.cuda() netD_B = netD_B.cuda() netG_A2B_optimizer = optimizer.Adam(params=netG_A2B.parameters(), lr=opt.lr, betas=(0.5, 0.999)) netG_B2A_optimizer = optimizer.Adam(params=netG_B2A.parameters(), lr=opt.lr, betas=(0.5, 0.999)) netD_A_optimizer = optimizer.Adam(params=netD_A.parameters(), lr=opt.lr, betas=(0.5, 0.999)) netD_B_optimizer = optimizer.Adam(params=netD_B.parameters(), lr=opt.lr, betas=(0.5, 0.999)) optimizers = dict() optimizers['G1'] = netG_A2B_optimizer optimizers['G2'] = netG_B2A_optimizer optimizers['D1'] = netD_A_optimizer optimizers['D2'] = netD_B_optimizer # Dataset loader transforms_ = [ transforms.RandomHorizontalFlip(), transforms.ToTensor(), transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)) ] tarindataloader = DataLoader(ImageDataset(opt.dataroot, transforms_=transforms_, unaligned=True), batch_size=opt.batchSize, shuffle=True) #writer writer = SummaryWriter(opt.log_dir) for epoch in range(0, opt.n_epochs): for ii, batch in enumerate(tarindataloader): # Set model input real_A = Variable(batch['A']) real_B = Variable(batch['B']) if opt.use_cuda: real_A = real_A.cuda() real_B = real_B.cuda() train_one_step(use_cuda=opt.use_cuda, netG_A2B=netG_A2B, netG_B2A=netG_B2A, netD_A=netD_A, netD_B=netD_B, real_A=real_A, real_B=real_B, optimizers=optimizers, iteration=ii, writer=writer) print("\nEpoch: %s Batch: %s" % (epoch, ii)) writer.export_scalars_to_json("./all_scalars.json") writer.close() torch.save(netG_A2B.state_dict(), os.path.join(opt.save_dir, '%s' % "netG_A2B")) torch.save(netG_B2A.state_dict(), os.path.join(opt.save_dir, '%s' % "netG_B2A")) torch.save(netD_A.state_dict(), os.path.join(opt.save_dir, '%s' % "netD_A")) torch.save(netD_B.state_dict(), os.path.join(opt.save_dir, '%s' % "netD_B"))
x = x.to(device) G_recon, _ = netG(x) result = torch.cat((x[0], G_recon[0]), 2) path = os.path.join( name + '_results', 'Transfer', str(epoch + 1) + '_epoch_' + name + '_test_' + str(n + 1) + '.png') plt.imsave(path, (result.cpu().numpy().transpose(1, 2, 0) + 1) / 2) if n == 4: break torch.save(netG.state_dict(), os.path.join(name + '_results', 'generator_latest.pkl')) torch.save( netD.state_dict(), os.path.join(name + '_results', 'discriminator_latest.pkl')) total_time = time.time() - start_time train_hist['total_time'].append(total_time) print("Avg one epoch time: %.2f, total %d epochs time: %.2f" % (torch.mean( torch.FloatTensor(train_hist['per_epoch_time'])), num_epochs, total_time)) print("Training finish!... save training results") torch.save(netG.state_dict(), os.path.join(name + '_results', 'generator_param.pkl')) torch.save(netD.state_dict(), os.path.join(name + '_results', 'discriminator_param.pkl')) with open(os.path.join(name + '_results', 'train_hist.pkl'), 'wb') as f: pickle.dump(train_hist, f)
output = model_d(g_out, onehotv) errG = criterion(output, labelv) optim_g.zero_grad() errG.backward() optim_g.step() d_loss += errD.data[0] g_loss += errG.data[0] if batch_idx % args.print_every == 0: print( "\t{} ({} / {}) mean D(fake) = {:.4f}, mean D(real) = {:.4f}". format(epoch_idx, batch_idx, len(train_loader), fakeD_mean, realD_mean)) g_out = model_g(fixed_noise, fixed_labels).data.view( SAMPLE_SIZE, 1, 28,28).cpu() save_image(g_out, '{}/{}_{}.png'.format( args.samples_dir, epoch_idx, batch_idx)) print('Epoch {} - D loss = {:.4f}, G loss = {:.4f}'.format(epoch_idx, d_loss, g_loss)) if epoch_idx % args.save_every == 0: torch.save({'state_dict': model_d.state_dict()}, '{}/model_d_epoch_{}.pth'.format( args.save_dir, epoch_idx)) torch.save({'state_dict': model_g.state_dict()}, '{}/model_g_epoch_{}.pth'.format( args.save_dir, epoch_idx))
class trainer(object): def __init__(self, cfg): self.cfg = cfg self.OldLabel_generator = U_Net(in_ch=cfg.DATASET.N_CLASS, out_ch=cfg.DATASET.N_CLASS, side='out') self.Image_generator = U_Net(in_ch=3, out_ch=cfg.DATASET.N_CLASS, side='in') self.discriminator = Discriminator(cfg.DATASET.N_CLASS + 3, cfg.DATASET.IMGSIZE, patch=True) self.criterion_G = GeneratorLoss(cfg.LOSS.LOSS_WEIGHT[0], cfg.LOSS.LOSS_WEIGHT[1], cfg.LOSS.LOSS_WEIGHT[2], ignore_index=cfg.LOSS.IGNORE_INDEX) self.criterion_D = DiscriminatorLoss() train_dataset = BaseDataset(cfg, split='train') valid_dataset = BaseDataset(cfg, split='val') self.train_dataloader = data.DataLoader( train_dataset, batch_size=cfg.DATASET.BATCHSIZE, num_workers=8, shuffle=True, drop_last=True) self.valid_dataloader = data.DataLoader( valid_dataset, batch_size=cfg.DATASET.BATCHSIZE, num_workers=8, shuffle=True, drop_last=True) self.ckpt_outdir = os.path.join(cfg.TRAIN.OUTDIR, 'checkpoints') if not os.path.isdir(self.ckpt_outdir): os.mkdir(self.ckpt_outdir) self.val_outdir = os.path.join(cfg.TRAIN.OUTDIR, 'val') if not os.path.isdir(self.val_outdir): os.mkdir(self.val_outdir) self.start_epoch = cfg.TRAIN.RESUME self.n_epoch = cfg.TRAIN.N_EPOCH self.optimizer_G = torch.optim.Adam( [{ 'params': self.OldLabel_generator.parameters() }, { 'params': self.Image_generator.parameters() }], lr=cfg.OPTIMIZER.G_LR, betas=(cfg.OPTIMIZER.BETA1, cfg.OPTIMIZER.BETA2), # betas=(cfg.OPTIMIZER.BETA1, cfg.OPTIMIZER.BETA2), weight_decay=cfg.OPTIMIZER.WEIGHT_DECAY) self.optimizer_D = torch.optim.Adam( [{ 'params': self.discriminator.parameters(), 'initial_lr': cfg.OPTIMIZER.D_LR }], lr=cfg.OPTIMIZER.D_LR, betas=(cfg.OPTIMIZER.BETA1, cfg.OPTIMIZER.BETA2), # betas=(cfg.OPTIMIZER.BETA1, cfg.OPTIMIZER.BETA2), weight_decay=cfg.OPTIMIZER.WEIGHT_DECAY) iter_per_epoch = len(train_dataset) // cfg.DATASET.BATCHSIZE lambda_poly = lambda iters: pow( (1.0 - iters / (cfg.TRAIN.N_EPOCH * iter_per_epoch)), 0.9) self.scheduler_G = torch.optim.lr_scheduler.LambdaLR( self.optimizer_G, lr_lambda=lambda_poly, ) # last_epoch=(self.start_epoch+1)*iter_per_epoch) self.scheduler_D = torch.optim.lr_scheduler.LambdaLR( self.optimizer_D, lr_lambda=lambda_poly, ) # last_epoch=(self.start_epoch+1)*iter_per_epoch) self.logger = logger(cfg.TRAIN.OUTDIR, name='train') self.running_metrics = runningScore(n_classes=cfg.DATASET.N_CLASS) if self.start_epoch >= 0: self.OldLabel_generator.load_state_dict( torch.load( os.path.join(cfg.TRAIN.OUTDIR, 'checkpoints', '{}epoch.pth'.format( self.start_epoch)))['model_G_N']) self.Image_generator.load_state_dict( torch.load( os.path.join(cfg.TRAIN.OUTDIR, 'checkpoints', '{}epoch.pth'.format( self.start_epoch)))['model_G_I']) self.discriminator.load_state_dict( torch.load( os.path.join(cfg.TRAIN.OUTDIR, 'checkpoints', '{}epoch.pth'.format( self.start_epoch)))['model_D']) self.optimizer_G.load_state_dict( torch.load( os.path.join(cfg.TRAIN.OUTDIR, 'checkpoints', '{}epoch.pth'.format( self.start_epoch)))['optimizer_G']) self.optimizer_D.load_state_dict( torch.load( os.path.join(cfg.TRAIN.OUTDIR, 'checkpoints', '{}epoch.pth'.format( self.start_epoch)))['optimizer_D']) log = "Using the {}th checkpoint".format(self.start_epoch) self.logger.info(log) self.Image_generator = self.Image_generator.cuda() self.OldLabel_generator = self.OldLabel_generator.cuda() self.discriminator = self.discriminator.cuda() self.criterion_G = self.criterion_G.cuda() self.criterion_D = self.criterion_D.cuda() def train(self): all_train_iter_total_loss = [] all_train_iter_corr_loss = [] all_train_iter_recover_loss = [] all_train_iter_change_loss = [] all_train_iter_gan_loss_gen = [] all_train_iter_gan_loss_dis = [] all_val_epo_iou = [] all_val_epo_acc = [] iter_num = [0] epoch_num = [] num_batches = len(self.train_dataloader) for epoch_i in range(self.start_epoch + 1, self.n_epoch): iter_total_loss = AverageTracker() iter_corr_loss = AverageTracker() iter_recover_loss = AverageTracker() iter_change_loss = AverageTracker() iter_gan_loss_gen = AverageTracker() iter_gan_loss_dis = AverageTracker() batch_time = AverageTracker() tic = time.time() # train self.OldLabel_generator.train() self.Image_generator.train() self.discriminator.train() for i, meta in enumerate(self.train_dataloader): image, old_label, new_label = meta[0].cuda(), meta[1].cuda( ), meta[2].cuda() recover_pred, feats = self.OldLabel_generator( label2onehot(old_label, self.cfg.DATASET.N_CLASS)) corr_pred = self.Image_generator(image, feats) # ------------------- # Train Discriminator # ------------------- self.discriminator.set_requires_grad(True) self.optimizer_D.zero_grad() fake_sample = torch.cat((image, corr_pred), 1).detach() real_sample = torch.cat( (image, label2onehot(new_label, cfg.DATASET.N_CLASS)), 1) score_fake_d = self.discriminator(fake_sample) score_real = self.discriminator(real_sample) gan_loss_dis = self.criterion_D(pred_score=score_fake_d, real_score=score_real) gan_loss_dis.backward() self.optimizer_D.step() self.scheduler_D.step() # --------------- # Train Generator # --------------- self.discriminator.set_requires_grad(False) self.optimizer_G.zero_grad() score_fake = self.discriminator( torch.cat((image, corr_pred), 1)) total_loss, corr_loss, recover_loss, change_loss, gan_loss_gen = self.criterion_G( corr_pred, recover_pred, score_fake, old_label, new_label) total_loss.backward() self.optimizer_G.step() self.scheduler_G.step() iter_total_loss.update(total_loss.item()) iter_corr_loss.update(corr_loss.item()) iter_recover_loss.update(recover_loss.item()) iter_change_loss.update(change_loss.item()) iter_gan_loss_gen.update(gan_loss_gen.item()) iter_gan_loss_dis.update(gan_loss_dis.item()) batch_time.update(time.time() - tic) tic = time.time() log = '{}: Epoch: [{}][{}/{}], Time: {:.2f}, ' \ 'Total Loss: {:.6f}, Corr Loss: {:.6f}, Recover Loss: {:.6f}, Change Loss: {:.6f}, GAN_G Loss: {:.6f}, GAN_D Loss: {:.6f}'.format( datetime.now(), epoch_i, i, num_batches, batch_time.avg, total_loss.item(), corr_loss.item(), recover_loss.item(), change_loss.item(), gan_loss_gen.item(), gan_loss_dis.item()) print(log) if (i + 1) % 10 == 0: all_train_iter_total_loss.append(iter_total_loss.avg) all_train_iter_corr_loss.append(iter_corr_loss.avg) all_train_iter_recover_loss.append(iter_recover_loss.avg) all_train_iter_change_loss.append(iter_change_loss.avg) all_train_iter_gan_loss_gen.append(iter_gan_loss_gen.avg) all_train_iter_gan_loss_dis.append(iter_gan_loss_dis.avg) iter_total_loss.reset() iter_corr_loss.reset() iter_recover_loss.reset() iter_change_loss.reset() iter_gan_loss_gen.reset() iter_gan_loss_dis.reset() vis.line(X=np.column_stack( np.repeat(np.expand_dims(iter_num, 0), 6, axis=0)), Y=np.column_stack((all_train_iter_total_loss, all_train_iter_corr_loss, all_train_iter_recover_loss, all_train_iter_change_loss, all_train_iter_gan_loss_gen, all_train_iter_gan_loss_dis)), opts={ 'legend': [ 'total_loss', 'corr_loss', 'recover_loss', 'change_loss', 'gan_loss_gen', 'gan_loss_dis' ], 'linecolor': np.array([[255, 0, 0], [0, 255, 0], [0, 0, 255], [255, 255, 0], [0, 255, 255], [255, 0, 255]]), 'title': 'Train loss of generator and discriminator' }, win='Train loss of generator and discriminator') iter_num.append(iter_num[-1] + 1) # eval self.OldLabel_generator.eval() self.Image_generator.eval() self.discriminator.eval() with torch.no_grad(): for j, meta in enumerate(self.valid_dataloader): image, old_label, new_label = meta[0].cuda(), meta[1].cuda( ), meta[2].cuda() recover_pred, feats = self.OldLabel_generator( label2onehot(old_label, self.cfg.DATASET.N_CLASS)) corr_pred = self.Image_generator(image, feats) preds = np.argmax(corr_pred.cpu().detach().numpy().copy(), axis=1) target = new_label.cpu().detach().numpy().copy() self.running_metrics.update(target, preds) if j == 0: color_map1 = gen_color_map(preds[0, :]).astype( np.uint8) color_map2 = gen_color_map(preds[1, :]).astype( np.uint8) color_map = cv2.hconcat([color_map1, color_map2]) cv2.imwrite( os.path.join( self.val_outdir, '{}epoch*{}*{}.png'.format( epoch_i, meta[3][0], meta[3][1])), color_map) score = self.running_metrics.get_scores() oa = score['Overall Acc: \t'] precision = score['Precision: \t'][1] recall = score['Recall: \t'][1] iou = score['Class IoU: \t'][1] miou = score['Mean IoU: \t'] self.running_metrics.reset() epoch_num.append(epoch_i) all_val_epo_acc.append(oa) all_val_epo_iou.append(miou) vis.line(X=np.column_stack( np.repeat(np.expand_dims(epoch_num, 0), 2, axis=0)), Y=np.column_stack((all_val_epo_acc, all_val_epo_iou)), opts={ 'legend': ['val epoch Overall Acc', 'val epoch Mean IoU'], 'linecolor': np.array([[255, 0, 0], [0, 255, 0]]), 'title': 'Validate Accuracy and IoU' }, win='validate Accuracy and IoU') log = '{}: Epoch Val: [{}], ACC: {:.2f}, Recall: {:.2f}, mIoU: {:.4f}' \ .format(datetime.now(), epoch_i, oa, recall, miou) self.logger.info(log) state = { 'epoch': epoch_i, "acc": oa, "recall": recall, "iou": miou, 'model_G_N': self.OldLabel_generator.state_dict(), 'model_G_I': self.Image_generator.state_dict(), 'model_D': self.discriminator.state_dict(), 'optimizer_G': self.optimizer_G.state_dict(), 'optimizer_D': self.optimizer_D.state_dict() } save_path = os.path.join(self.cfg.TRAIN.OUTDIR, 'checkpoints', '{}epoch.pth'.format(epoch_i)) torch.save(state, save_path)
class GAIL(PPO): def __init__( self, state_dimension: Tuple, action_space: int, save_path: Path, hyp: HyperparametersGAIL, policy_params: namedtuple, discriminator_params: DiscrimParams, param_plot_num: int, ppo_type: str = "clip", adv_type: str = "monte_carlo", max_plot_size: int = 10000, policy_burn_in: int = 0, verbose: bool = False, ): self.discrim_net_save = save_path / "GAIL_discrim.pth" self.discriminator = Discriminator( state_dimension, action_space, discriminator_params, ).to(device) self.discrim_optim = torch.optim.Adam(self.discriminator.parameters(), lr=hyp.discrim_lr) gail_plots = [("discrim_loss", np.float64)] super(GAIL, self).__init__( state_dimension, action_space, save_path, hyp, policy_params, param_plot_num, ppo_type, advantage_type=adv_type, neural_net_save=f"GAIL-{adv_type}", max_plot_size=max_plot_size, discrim_params=discriminator_params, policy_burn_in=policy_burn_in, verbose=verbose, additional_plots=gail_plots, ) self.discrim_loss = torch.nn.NLLLoss() def update(self, buffer: GAILExperienceBuffer, ep_num: int): # Update discriminator state_actions = buffer.state_actions.to(device) num_learner_samples = buffer.get_length() expert_samples_per_epoch = int( (state_actions.size()[0] - num_learner_samples) / self.hyp.num_discrim_epochs) for epoch in range(self.hyp.num_discrim_epochs): step_state_actions = torch.cat( ( state_actions[:num_learner_samples], state_actions[ num_learner_samples + epoch * expert_samples_per_epoch:num_learner_samples + (epoch + 1) * expert_samples_per_epoch], ), dim=0, ) discrim_logprobs = self.discriminator.logprobs( step_state_actions).to(device) loss = self.discrim_loss( input=discrim_logprobs, target=buffer.discrim_labels.type(torch.long), ) plotted_loss = loss.detach().cpu().numpy() self.plotter.record_data({"discrim_loss": plotted_loss}) if self.verbose: print( f"Learner labels {buffer.discrim_labels[:num_learner_samples].mean()}: " f"\t{torch.exp(discrim_logprobs[:num_learner_samples]).t()[1].mean()}" ) print( f"Expert labels {buffer.discrim_labels[num_learner_samples:].mean()}: " f"\t\t{torch.exp(discrim_logprobs[num_learner_samples:]).t()[1].mean()}" ) self.discrim_optim.zero_grad() loss.backward() self.discrim_optim.step() self.record_nn_params() # Update policy buffer.rewards = list( np.squeeze( self.discriminator.logprob_expert( state_actions[:num_learner_samples]).float().detach().cpu( ).numpy())) if self.verbose: print( "----------------------------------------------------------------------" ) super(GAIL, self).update(buffer, ep_num) def record_nn_params(self): """Gets randomly sampled actor NN parameters from 1st layer.""" names, x_params, y_params = self.plotter.get_param_plot_nums() sampled_params = {} for name, x_param, y_param in zip(names, x_params, y_params): network_to_sample = (self.discriminator if name[:7] == "discrim" else self.policy) sampled_params[name] = ( network_to_sample.state_dict()[name].cpu().numpy()[x_param, y_param]) self.plotter.record_data(sampled_params) def _save_network(self): super(GAIL, self)._save_network() torch.save(self.discriminator.state_dict(), f"{self.discrim_net_save}") def _load_network(self): super(GAIL, self)._load_network() print( f"Loading discriminator network saved at: {self.discrim_net_save}") net = torch.load(self.discrim_net_save, map_location=device) self.discriminator.load_state_dict(net)
class CycleGAN(AlignmentModel): """This class implements the alignment model for GAN networks with two generators and two discriminators (cycle GAN). For description of the implemented functions, refer to the alignment model.""" def __init__(self, device, config, generator_a=None, generator_b=None, discriminator_a=None, discriminator_b=None): """Initialize two new generators and two discriminators from the config or use pre-trained ones and create Adam optimizers for all models.""" super().__init__(device, config) self.epoch_losses = [0., 0., 0., 0.] if generator_a is None: generator_a_conf = dict( dim_1=config['dim_b'], dim_2=config['dim_a'], layer_number=config['generator_layers'], layer_expansion=config['generator_expansion'], initialize_generator=config['initialize_generator'], norm=config['gen_norm'], batch_norm=config['gen_batch_norm'], activation=config['gen_activation'], dropout=config['gen_dropout']) self.generator_a = Generator(generator_a_conf, device) self.generator_a.to(device) else: self.generator_a = generator_a if 'optimizer' in config: self.optimizer_g_a = OPTIMIZERS[config['optimizer']]( self.generator_a.parameters(), config['learning_rate']) elif 'optimizer_default' in config: if config['optimizer_default'] == 'sgd': self.optimizer_g_a = OPTIMIZERS[config['optimizer_default']]( self.generator_a.parameters(), config['learning_rate']) else: self.optimizer_g_a = OPTIMIZERS[config['optimizer_default']]( self.generator_a.parameters()) else: self.optimizer_g_a = torch.optim.Adam( self.generator_a.parameters(), config['learning_rate']) if generator_b is None: generator_b_conf = dict( dim_1=config['dim_a'], dim_2=config['dim_b'], layer_number=config['generator_layers'], layer_expansion=config['generator_expansion'], initialize_generator=config['initialize_generator'], norm=config['gen_norm'], batch_norm=config['gen_batch_norm'], activation=config['gen_activation'], dropout=config['gen_dropout']) self.generator_b = Generator(generator_b_conf, device) self.generator_b.to(device) else: self.generator_b = generator_b if 'optimizer' in config: self.optimizer_g_b = OPTIMIZERS[config['optimizer']]( self.generator_b.parameters(), config['learning_rate']) elif 'optimizer_default' in config: if config['optimizer_default'] == 'sgd': self.optimizer_g_b = OPTIMIZERS[config['optimizer_default']]( self.generator_b.parameters(), config['learning_rate']) else: self.optimizer_g_b = OPTIMIZERS[config['optimizer_default']]( self.generator_b.parameters()) else: self.optimizer_g_b = torch.optim.Adam( self.generator_b.parameters(), config['learning_rate']) if discriminator_a is None: discriminator_a_conf = dict( dim=config['dim_a'], layer_number=config['discriminator_layers'], layer_expansion=config['discriminator_expansion'], batch_norm=config['disc_batch_norm'], activation=config['disc_activation'], dropout=config['disc_dropout']) self.discriminator_a = Discriminator(discriminator_a_conf, device) self.discriminator_a.to(device) else: self.discriminator_a = discriminator_a if 'optimizer' in config: self.optimizer_d_a = OPTIMIZERS[config['optimizer']]( self.discriminator_a.parameters(), config['learning_rate']) elif 'optimizer_default' in config: if config['optimizer_default'] == 'sgd': self.optimizer_d_a = OPTIMIZERS[config['optimizer_default']]( self.discriminator_a.parameters(), config['learning_rate']) else: self.optimizer_d_a = OPTIMIZERS[config['optimizer_default']]( self.discriminator_a.parameters()) else: self.optimizer_d_a = torch.optim.Adam( self.discriminator_a.parameters(), config['learning_rate']) if discriminator_b is None: discriminator_b_conf = dict( dim=config['dim_b'], layer_number=config['discriminator_layers'], layer_expansion=config['discriminator_expansion'], batch_norm=config['disc_batch_norm'], activation=config['disc_activation'], dropout=config['disc_dropout']) self.discriminator_b = Discriminator(discriminator_b_conf, device) self.discriminator_b.to(device) else: self.discriminator_b = discriminator_b if 'optimizer' in config: self.optimizer_d_b = OPTIMIZERS[config['optimizer']]( self.discriminator_b.parameters(), config['learning_rate']) elif 'optimizer_default' in config: if config['optimizer_default'] == 'sgd': self.optimizer_d_b = OPTIMIZERS[config['optimizer_default']]( self.discriminator_b.parameters(), config['learning_rate']) else: self.optimizer_d_b = OPTIMIZERS[config['optimizer_default']]( self.discriminator_b.parameters()) else: self.optimizer_d_b = torch.optim.Adam( self.discriminator_b.parameters(), config['learning_rate']) def train(self): self.generator_a.train() self.generator_b.train() self.discriminator_a.train() self.discriminator_b.train() def eval(self): self.generator_a.eval() self.generator_b.eval() self.discriminator_a.eval() self.discriminator_b.eval() def zero_grad(self): self.optimizer_g_a.zero_grad() self.optimizer_g_b.zero_grad() self.optimizer_d_a.zero_grad() self.optimizer_d_b.zero_grad() def optimize_all(self): self.optimizer_g_a.step() self.optimizer_g_b.step() self.optimizer_d_a.step() self.optimizer_d_b.step() def optimize_generator(self): """Do the optimization step only for generators (e.g. when training generators and discriminators separately or in turns).""" self.optimizer_g_a.step() self.optimizer_g_b.step() def optimize_discriminator(self): """Do the optimization step only for discriminators (e.g. when training generators and discriminators separately or in turns).""" self.optimizer_d_a.step() self.optimizer_d_b.step() def change_lr(self, factor): self.current_lr = self.current_lr * factor for param_group in self.optimizer_g_a.param_groups: param_group['lr'] = self.current_lr for param_group in self.optimizer_g_b.param_groups: param_group['lr'] = self.current_lr def update_losses_batch(self, *losses): loss_g_a, loss_g_b, loss_d_a, loss_d_b = losses self.epoch_losses[0] += loss_g_a self.epoch_losses[1] += loss_g_b self.epoch_losses[2] += loss_d_a self.epoch_losses[3] += loss_d_b def complete_epoch(self, epoch_metrics): self.metrics.append(epoch_metrics + [sum(self.epoch_losses)]) self.losses.append(self.epoch_losses) self.epoch_losses = [0., 0., 0., 0.] def print_epoch_info(self): print( f"{len(self.metrics)} ### {self.losses[-1][0]:.2f} - {self.losses[-1][1]:.2f} " f"- {self.losses[-1][2]:.2f} - {self.losses[-1][3]:.2f} ### {self.metrics[-1]}" ) def copy_model(self): self.model_copy = deepcopy(self.generator_a.state_dict()), deepcopy(self.generator_b.state_dict()),\ deepcopy(self.discriminator_a.state_dict()), deepcopy(self.discriminator_b.state_dict()) def restore_model(self): self.generator_a.load_state_dict(self.model_copy[0]) self.generator_b.load_state_dict(self.model_copy[1]) self.discriminator_a.load_state_dict(self.model_copy[2]) self.discriminator_b.load_state_dict(self.model_copy[3]) def export_model(self, test_results, description=None): if description is None: description = f"CycleGAN_{self.config['evaluation']}_{self.config['subset']}" export_cyclegan_alignment(description, self.config, self.generator_a, self.generator_b, self.discriminator_a, self.discriminator_b, self.metrics) save_alignment_test_results(test_results, description) print(f"Saved model to directory {description}.") @classmethod def load_model(cls, name, device): generator_a, generator_b, discriminator_a, discriminator_b, config = load_cyclegan_alignment( name, device) model = cls(device, config, generator_a, generator_b, discriminator_a, discriminator_b) return model
class GAIL: def __init__(self, exp_dir, exp_thresh, state_dim, action_dim, learn_rate, betas, _device, _gamma, load_weights=False): """ exp_dir : directory containing the expert episodes exp_thresh : parameter to control number of episodes to load as expert based on returns (lower means more episodes) state_dim : dimesnion of state action_dim : dimesnion of action learn_rate : learning rate for optimizer _device : GPU or cpu _gamma : discount factor _load_weights : load weights from directory """ # storing runtime device self.device = _device # discount factor self.gamma = _gamma # Expert trajectory self.expert = ExpertTrajectories(exp_dir, exp_thresh, gamma=self.gamma) # Defining the actor and its optimizer self.actor = ActorNetwork(state_dim).to(self.device) self.optim_actor = torch.optim.Adam(self.actor.parameters(), lr=learn_rate, betas=betas) # Defining the discriminator and its optimizer self.disc = Discriminator(state_dim, action_dim).to(self.device) self.optim_disc = torch.optim.Adam(self.disc.parameters(), lr=learn_rate, betas=betas) if not load_weights: self.actor.apply(init_weights) self.disc.apply(init_weights) else: self.load() # Loss function crtiterion self.criterion = torch.nn.BCELoss() def get_action(self, state): """ obtain action for a given state using actor network """ state = torch.tensor(state, dtype=torch.float, device=self.device).view(1, -1) return self.actor(state).cpu().data.numpy().flatten() def update(self, n_iter, batch_size=100): """ train discriminator and actor for mini-batch """ # memory to store disc_losses = np.zeros(n_iter, dtype=np.float) act_losses = np.zeros(n_iter, dtype=np.float) for i in range(n_iter): # Get expert state and actions batch exp_states, exp_actions = self.expert.sample(batch_size) exp_states = torch.FloatTensor(exp_states).to(self.device) exp_actions = torch.FloatTensor(exp_actions).to(self.device) # Get state, and actions using actor states, _ = self.expert.sample(batch_size) states = torch.FloatTensor(states).to(self.device) actions = self.actor(states) ''' train the discriminator ''' self.optim_disc.zero_grad() # label tensors exp_labels = torch.full((batch_size, 1), 1, device=self.device) policy_labels = torch.full((batch_size, 1), 0, device=self.device) # with expert transitions prob_exp = self.disc(exp_states, exp_actions) exp_loss = self.criterion(prob_exp, exp_labels) # with policy actor transitions prob_policy = self.disc(states, actions.detach()) policy_loss = self.criterion(prob_policy, policy_labels) # use backprop disc_loss = exp_loss + policy_loss disc_losses[i] = disc_loss.mean().item() disc_loss.backward() self.optim_disc.step() ''' train the actor ''' self.optim_actor.zero_grad() loss_actor = -self.disc(states, actions) act_losses[i] = loss_actor.mean().detach().item() loss_actor.mean().backward() self.optim_actor.step() print("Finished training minibatch") return act_losses, disc_losses def save( self, directory='/home/aman/Programming/RL-Project/Deterministic-GAIL/weights', name='GAIL'): torch.save(self.actor.state_dict(), '{}/{}_actor.pth'.format(directory, name)) torch.save(self.disc.state_dict(), '{}/{}_discriminator.pth'.format(directory, name)) def load( self, directory='/home/aman/Programming/RL-Project/Deterministic-GAIL/weights', name='GAIL'): print(os.getcwd()) self.actor.load_state_dict( torch.load('{}/{}_actor.pth'.format(directory, name))) self.disc.load_state_dict( torch.load('{}/{}_discriminator.pth'.format(directory, name))) def set_mode(self, mode="train"): if mode == "train": self.actor.train() self.disc.train() else: self.actor.eval() self.disc.eval()
def train_gan(args): # prepare dataloader dataloader = create_data_loader(args) # set up device device = torch.device('cuda:0' if ( torch.cuda.is_available() and args.ngpu > 0) else 'cpu') # Create & setup generator netG = Generator(args).to(device) # handle multiple gpus if (device.type == 'cuda' and args.ngpu > 1): netG = nn.DataParallel(netG, list(range(args.ngpu))) # load from checkpoint if available if args.netG: netG.load_state_dict(torch.load(args.netG)) # initialize network with random weights else: netG.apply(weights_init) # Create & setup discriminator netD = Discriminator(args).to(device) # handle multiple gpus if (device.type == 'cuda' and args.ngpu > 1): netD = nn.DataParallel(netD, list(range(args.ngpu))) # load from checkpoint if available if args.netD: netD.load_state_dict(torch.load(args.netD)) # initialize network with random weights else: netD.apply(weights_init) # setup up loss & optimizers criterion = nn.BCELoss() optimizerG = optim.Adam(netG.parameters(), lr=args.lr, betas=(args.beta1, 0.999)) optimizerD = optim.Adam(netD.parameters(), lr=args.lr, betas=(args.beta1, 0.999)) # For input of generator in testing fixed_noise = torch.randn(64, args.nz, 1, 1, device=device) # convention for training real_label = 1 fake_label = 0 # training data for later analysis img_list = [] G_losses = [] D_losses = [] iters = 0 # epochs num_epochs = 150 print('Starting Training Loop....') # For each epoch for e in range(args.num_epochs): # for each batch in the dataloader for i, data in enumerate(dataloader, 0): ########## Training Discriminator ########## netD.zero_grad() # train with real data real_data = data[0].to(device) # make labels batch_size = real_data.size(0) labels = torch.full((batch_size, ), real_label, device=device) # forward pass real data through D real_outputD = netD(real_data).view(-1) # calc error on real data errD_real = criterion(real_outputD, labels) # calc grad errD_real.backward() D_x = real_outputD.mean().item() # train with fake data noise = torch.randn(batch_size, args.nz, 1, 1, device=device) fake_data = netG(noise) labels.fill_(fake_label) # classify fake fake_outputD = netD(fake_data.detach()).view(-1) # calc error on fake data errD_fake = criterion(fake_outputD, labels) # calc grad errD_fake.backward() D_G_z1 = fake_outputD.mean().item() # add all grad and update D errD = errD_real + errD_fake optimizerD.step() ######################################## ########## Training Generator ########## netG.zero_grad() # since aim is fooling the netD, labels should be flipped labels.fill_(real_label) # forward pass with updated netD fake_outputD = netD(fake_data).view(-1) # calc error errG = criterion(fake_outputD, labels) # calc grad errG.backward() D_G_z2 = fake_outputD.mean().item() # update G optimizerG.step() ######################################## # output training stats if i % 500 == 0: print(f'[{e+1}/{args.num_epochs}][{i+1}/{len(dataloader)}]\ \tLoss_D:{errD.item():.4f}\ \tLoss_G:{errG.item():.4f}\ \tD(x):{D_x:.4f}\ \tD(G(z)):{D_G_z1:.4f}/{D_G_z2:.4f}') # for later plot G_losses.append(errG.item()) D_losses.append(errD.item()) # generate fake image on fixed noise for comparison if ((iters % 500 == 0) or ((e == args.num_epochs - 1) and (i == len(dataloader) - 1))): with torch.no_grad(): fake = netG(fixed_noise).detach().cpu() img_list.append( vutils.make_grid(fake, padding=2, normalize=True)) iters += 1 if e % args.save_every == 0: # save at args.save_every epoch torch.save(netG.state_dict(), args.outputG) torch.save(netD.state_dict(), args.outputD) print(f'Made a New Checkpoint for {e+1}') # return training data for analysis return img_list, G_losses, D_losses