def __init__(self, args): kwargs = {'num_workers': 4, 'pin_memory': True} self.source_loader, self.target_loader, self.test_loader, self.nclass = make_data_loader(args, **kwargs) self.tbar = tqdm(self.test_loader, desc='\r') self.trainer = wgan_trainer(args, 2) self.evaluator = Evaluator(2) self.best_IoU = {'disc': 0.77, 'cup': 0.65} self.attempt = 9.5 self.validation(args, self.trainer.target_model, self.tbar ) self.trainer_wgan(args)
class wgan: def __init__(self, args): kwargs = {'num_workers': 4, 'pin_memory': True} self.source_loader, self.target_loader, self.test_loader, self.nclass = make_data_loader( args, **kwargs) self.tbar = tqdm(self.test_loader, desc='\r') self.trainer = wgan_trainer(args, 2) self.evaluator = Evaluator(2) self.best_IoU = {'disc': 0.77, 'cup': 0.65} self.attempt = 9.5 self.validation(args, self.trainer.target_model, self.tbar) self.trainer_wgan(args) def loop_iterable(self, iterable): while True: yield from iterable def save_model(self, epoch): print('Validation:') Acc = self.evaluator.Pixel_Accuracy([ self.evaluator.confusion_matrix_disc, self.evaluator.confusion_matrix_cup ]) Acc_class = self.evaluator.Pixel_Accuracy_Class([ self.evaluator.confusion_matrix_disc, self.evaluator.confusion_matrix_cup ]) mIoU = self.evaluator.Mean_Intersection_over_Union([ self.evaluator.confusion_matrix_disc, self.evaluator.confusion_matrix_cup ]) FWIoU = self.evaluator.Frequency_Weighted_Intersection_over_Union([ self.evaluator.confusion_matrix_disc, self.evaluator.confusion_matrix_cup ]) print("epoch:{}, Acc:{}, Acc_class:{}, mIoU:{}, fwIoU: {}".format( epoch, Acc, Acc_class, mIoU, FWIoU)) if (mIoU['cup'] > self.best_IoU['cup']): #model save self.best_IoU = mIoU print('---- MODEL SAVE ---') torch.save( { 'epoch': epoch + 1, 'state_dict': self.trainer.target_model.state_dict(), 'best_auc': str(mIoU['cup']), 'optimizer': self.trainer.dda_optim.state_dict() }, 'best_models_save/groupwise/m-adda_wgan_clip_0.03' + "v_" + str(self.attempt) + '.pth.tar') return mIoU def trainer_wgan(self, args): self.trainer.target_model.train() self.trainer.disc_model.train() self.evaluator.reset() max_epochs = args.epochs for epoch in range(1, max_epochs + 1): self.trainer.target_model.train() batch_iterator = zip(self.loop_iterable(self.source_loader), self.loop_iterable(self.target_loader)) total_loss = 0 total_loss_tgt = 0 loss_critic = 0 loss_tgt = 0 total_accuracy = 0 regu_val_disc = 0 len_dataloader = max(len(self.source_loader), len(self.target_loader)) torch.manual_seed(1 + epoch) for step in trange(len_dataloader, leave=True): p = float(step + epoch * len_dataloader ) / args.epochs / len_dataloader #(1+len)/epochs/ alpha = 2. / (1. + np.exp(-10 * p)) - 1 try: data = next(batch_iterator) except StopIteration: batch_iterator = zip( self.loop_iterable(self.source_loader), self.loop_iterable(self.target_loader)) data = next(batch_iterator) if epoch < 0: source_x, src_labels = data[0][0].cuda(), data[0][1].cuda() target_x, target_lab = data[1][0].cuda(), data[1][1].cuda() dda_loss, tgt_loss = self.trainer.update_weights( source_x, src_labels, target_x, target_lab, 0, 0, 'train_gen') continue for i in range(args.k_disc): source_x, src_labels = data[0][0].cuda(), data[0][1].cuda() target_x, target_lab = data[1][0].cuda(), data[1][1].cuda() # t = np.array(target_x.detach().cpu()[0,:,:,:]) # t = t.transpose((1,2,0)) # import cv2 # cv2.imwrite('cup_orig.png', t*255) # target_x, target_lab = data[1][0].cuda(), data[1][1].cuda() #target_x = self.trainer.adv_aug.perturb(target_x, target_lab, self.trainer.target_criterion, random_start=False ) # t = np.array(target_x.detach().cpu()[0,:,:,:]) # t = t.transpose((1,2,0)) # cv2.imwrite('cup.png', t*255) dda_loss, tgt_loss, regu_val_disc = self.trainer.update_weights( source_x, src_labels, target_x, target_lab, 0.1, 0.1, 0.8, 1.2, 0.01, 'train_disc') #0.2, 0.01 for i in range(args.k_src): source_x, src_labels = data[0][0].cuda(), data[0][1].cuda() target_x, target_lab = data[1][0].cuda(), data[1][1].cuda() #target_x = self.trainer.adv_aug.perturb(target_x, target_lab, self.trainer.target_criterion, random_start=False ) dda_loss, tgt_loss = self.trainer.update_weights( source_x, src_labels, target_x, target_lab, 0.1, 0.1, 0.8, 1.2, 0.01, 'train_gen') #0.2 , 0.01 total_loss += dda_loss total_loss_tgt += tgt_loss regu_val_disc += regu_val_disc if (step % 50 == 0): print(tgt_loss, dda_loss, regu_val_disc) # print("Target_loss:{}, disc_loss:{}".format(total_loss_tgt/(step+1), total_loss/(step+1))) self.trainer.scheduler(self.trainer.dda_optim, step, epoch, self.best_IoU['cup']) self.trainer.target_model.eval() self.trainer.disc_model.eval() for st, data in enumerate(self.tbar): image, target = data[0], data[1] #target = target.transpose(3,1) image, target = image.cuda(), target.cuda() with torch.no_grad(): output, _ = self.trainer.target_model(image) test_loss = self.trainer.target_criterion(output, target) pred = output.data.cpu().numpy() target = target.cpu().numpy() pred[pred >= 0.5] = 1 pred[pred < 0.5] = 0 # Add batch sample into evaluator self.evaluator.add_batch(target, pred) self.evaluator.add_test_loss(test_loss / (st + 1)) mIoU = self.evaluator.Mean_Intersection_over_Union([ self.evaluator.confusion_matrix_disc, self.evaluator.confusion_matrix_cup ]) print("mIoU:{}".format(mIoU)) total_accuracy = 0 if ((epoch + 1) % 1 == 0): print("Epoch [{}/{}] Step [{}/{}]:" "d_loss={:.5f} g_loss={:.5f} acc={:.5f}".format( epoch + 1, max_epochs, epoch + 1, len(self.source_loader), total_loss / ((step + 1)), total_loss_tgt / (step + 1), total_accuracy)) mIoU = self.save_model(epoch) def validation(self, args, model, tbar): best_pred = {'cup': 0, 'disc': 0} model.eval() self.evaluator.reset() test_loss = 0.0 for i, data in enumerate(tbar): image, target = data[0], data[1] #target = target.transpose(3,1) image, target = image.cuda(), target.cuda() import pdb with torch.no_grad(): output, _ = model(image) test_loss = self.trainer.target_criterion(output, target) pred = output.data.cpu().numpy() target = target.cpu().numpy() pred[pred >= 0.5] = 1 pred[pred < 0.5] = 0 # Add batch sample into evaluator self.evaluator.add_batch(target, pred) self.evaluator.add_test_loss(test_loss / (i + 1)) mIoU = self.evaluator.Mean_Intersection_over_Union([ self.evaluator.confusion_matrix_disc, self.evaluator.confusion_matrix_cup ]) #evaluator.Plot_Loss(1) print('Validation:') #print('[Epoch: %d, numImages: %5d]' % (epoch, i * args.batch_size + image.data.shape[0])) print("mIoU:{}".format(mIoU)) print('Loss: %.3f' % test_loss) new_pred = mIoU if new_pred['cup'] > best_pred['cup']: is_best = True best_pred = new_pred
def __init__(self, args): # Define Saver self.saver = Saver(args) self.saver.save_experiment_config() # Define Tensorboard Summary self.summary = TensorboardSummary(self.saver.experiment_dir) self.writer = self.summary.create_summary() checkpoint = torch.load(args.resume) args.start_epoch = checkpoint['epoch'] # Define Dataloader kwargs = {'num_workers': args.workers, 'pin_memory': True} self.source_loader, self.target_loader, _, self.nclass = make_data_loader( args, **kwargs) # Define Target Model self.target_model = DeepLab(num_classes=self.nclass, backbone=args.backbone, output_stride=args.out_stride, sync_bn=args.sync_bn, freeze_bn=args.freeze_bn) # Using cuda self.best_pred = {'disc': 0.0, 'cup': 0.0} self.target_model = torch.nn.DataParallel(self.target_model) patch_replication_callback(self.target_model) self.target_model = self.target_model.cuda() model_dict = self.target_model.module.state_dict() pretrained_dict = { k: v for k, v in checkpoint['state_dict'].items() if 'last_conv' not in k } model_dict.update(pretrained_dict) self.target_model.module.load_state_dict(model_dict) self.target_model.train() self.set_requires_grad('target', True) # Define learning rate and optimizer params target_params = [{ 'params': self.target_model.module.get_1x_lr_params(), 'lr': args.lr }, { 'params': self.target_model.module.get_10x_lr_params(), 'lr': args.lr * 10 }] target_optim = torch.optim.SGD(target_params, momentum=args.momentum, weight_decay=args.weight_decay, nesterov=args.nesterov) target_optim.zero_grad() self.target_criterion = torch.nn.BCEWithLogitsLoss() self.target_optim = target_optim # Define lr scheduler self.scheduler = LR_Scheduler(args.lr_scheduler, args.lr, args.epochs, len(self.target_loader)) self.evaluator = Evaluator(3)
class Adda(): def __init__(self, args): # Define Saver self.saver = Saver(args) self.saver.save_experiment_config() # Define Tensorboard Summary self.summary = TensorboardSummary(self.saver.experiment_dir) self.writer = self.summary.create_summary() checkpoint = torch.load(args.resume) args.start_epoch = checkpoint['epoch'] # Define Dataloader kwargs = {'num_workers': args.workers, 'pin_memory': True} self.source_loader, self.target_loader, _, self.nclass = make_data_loader( args, **kwargs) # Define Target Model self.target_model = DeepLab(num_classes=self.nclass, backbone=args.backbone, output_stride=args.out_stride, sync_bn=args.sync_bn, freeze_bn=args.freeze_bn) # Using cuda self.best_pred = {'disc': 0.0, 'cup': 0.0} self.target_model = torch.nn.DataParallel(self.target_model) patch_replication_callback(self.target_model) self.target_model = self.target_model.cuda() model_dict = self.target_model.module.state_dict() pretrained_dict = { k: v for k, v in checkpoint['state_dict'].items() if 'last_conv' not in k } model_dict.update(pretrained_dict) self.target_model.module.load_state_dict(model_dict) self.target_model.train() self.set_requires_grad('target', True) # Define learning rate and optimizer params target_params = [{ 'params': self.target_model.module.get_1x_lr_params(), 'lr': args.lr }, { 'params': self.target_model.module.get_10x_lr_params(), 'lr': args.lr * 10 }] target_optim = torch.optim.SGD(target_params, momentum=args.momentum, weight_decay=args.weight_decay, nesterov=args.nesterov) target_optim.zero_grad() self.target_criterion = torch.nn.BCEWithLogitsLoss() self.target_optim = target_optim # Define lr scheduler self.scheduler = LR_Scheduler(args.lr_scheduler, args.lr, args.epochs, len(self.target_loader)) self.evaluator = Evaluator(3) def set_requires_grad(self, mode, requires_grad=False): for param in eval('self.' + mode + '_model').parameters(): param.requires_grad = requires_grad def loop_iterable(self, iterable): while True: yield from iterable def trainer(self, num_epochs, iterations): best_IoU = 0.68 attempt = 1 for epoch in range(1, num_epochs + 1): total_loss_tgt = 0 total_accuracy = 0 self.evaluator.reset() total_loss = 0 len_dataloader = len(self.source_loader) torch.manual_seed(1 + epoch) for step in trange(len_dataloader, leave=True): p = float(step + epoch * len_dataloader ) / args.epochs / len_dataloader #(1+len)/epochs/ alpha = 2. / (1. + np.exp(-10 * p)) - 1 try: data = next(batch_iterator) except StopIteration: batch_iterator = zip( self.loop_iterable(self.source_loader), self.loop_iterable(self.target_loader)) data = next(batch_iterator)