def ranking_target_w_discrim(model_seg, target_loader, output_path, cfg): train_lst_ordered = osp.join(output_path,'train_ent_full.txt') writer = SummaryWriter(log_dir=cfg.TEST.SNAPSHOT_DIR) num_classes = cfg.NUM_CLASSES device = cfg.GPU_ID input_size_target = cfg.TRAIN.INPUT_SIZE_TARGET targetloader_iter = enumerate(target_loader) interp_target = nn.Upsample(size=(input_size_target[1], input_size_target[0]), mode='bilinear', align_corners=True) dis_dictionary = {} l1loss = torch.nn.L1Loss() with torch.no_grad(): for i in tqdm(range(len(target_loader))): # for i in tqdm(range(30)): _, batch = targetloader_iter.__next__() images, image_aug, _, _, name, _ = batch pred_trg_main, _ = model_seg(images.cuda(device)) # edge = CannyFilter(images) # pred_trg_main_aug, _ = model_seg(image_aug.cuda(device)) pred_trg_main = interp_target(pred_trg_main) # pred_trg_main_aug = F.softmax(interp_target(pred_trg_main_aug),dim=1) # pred_trg_max, _ = torch.max(pred_trg_main,dim=1) # pred_trg_main_aug_ = pred_trg_main * pred_trg_main_aug # diff = l1loss(pred_trg_main, pred_trg_main_aug_) # diff = torch.mean(diff) # pred_trg_entropy = model_dis(F.softmax(pred_trg_main)) pred_trg_entropy = prob_2_entropy(F.softmax(pred_trg_main)) # pred_trg_entropy = abs(float(pred_trg_entropy.mean()) - 0.5) # import pdb # pdb.set_trace() pred_trg_entropy = float(pred_trg_entropy.mean()) dis_dictionary[name[0]] = pred_trg_entropy # if i % 2 == 0: # draw_in_tensorboard(writer, images, image_aug, i, pred_trg_main, pred_trg_main_aug_, 'S') # import pdb # pdb.set_trace() dis_dictionary = sorted(dis_dictionary.items(),key=f2) with open(train_lst_ordered, 'w') as f: for i in range(len(dis_dictionary)): f.write("%s\t%s\n" % (dis_dictionary[i][0], dis_dictionary[i][1]))
def train_advent(model, trainloader, targetloader, cfg): ''' UDA training with advent ''' # Create the model and start the training. # pdb.set_trace() input_size_source = cfg.TRAIN.INPUT_SIZE_SOURCE input_size_target = cfg.TRAIN.INPUT_SIZE_TARGET device = cfg.GPU_ID num_classes = cfg.NUM_CLASSES viz_tensorboard = os.path.exists(cfg.TRAIN.TENSORBOARD_LOGDIR) if viz_tensorboard: writer = SummaryWriter(log_dir=cfg.TRAIN.TENSORBOARD_LOGDIR) # SEGMNETATION NETWORK model.train() model.to(device) cudnn.benchmark = True cudnn.enabled = True # DISCRIMINATOR NETWORK # feature-level d_aux = get_fc_discriminator(num_classes=num_classes) d_aux.train() d_aux.to(device) # restore_from = cfg.TRAIN.RESTORE_FROM_aux # print("Load Discriminator:", restore_from) # load_checkpoint_for_evaluation(d_aux, restore_from, device) # seg maps, i.e. output, level d_main = get_fc_discriminator(num_classes=num_classes) d_main.train() d_main.to(device) # restore_from = cfg.TRAIN.RESTORE_FROM_main # print("Load Discriminator:", restore_from) # load_checkpoint_for_evaluation(d_main, restore_from, device) # OPTIMIZERS # segnet's optimizer optimizer = optim.SGD(model.optim_parameters(cfg.TRAIN.LEARNING_RATE), lr=cfg.TRAIN.LEARNING_RATE, momentum=cfg.TRAIN.MOMENTUM, weight_decay=cfg.TRAIN.WEIGHT_DECAY) # discriminators' optimizers optimizer_d_aux = optim.Adam(d_aux.parameters(), lr=cfg.TRAIN.LEARNING_RATE_D, betas=(0.9, 0.99)) optimizer_d_main = optim.Adam(d_main.parameters(), lr=cfg.TRAIN.LEARNING_RATE_D, betas=(0.9, 0.99)) # interpolate output segmaps interp = nn.Upsample(size=(input_size_source[1], input_size_source[0]), mode='bilinear', align_corners=True) interp_target = nn.Upsample(size=(input_size_target[1], input_size_target[0]), mode='bilinear', align_corners=True) # labels for adversarial training source_label = 0 target_label = 1 trainloader_iter = enumerate(trainloader) targetloader_iter = enumerate(targetloader) for i_iter in tqdm(range(cfg.TRAIN.EARLY_STOP)): # reset optimizers optimizer.zero_grad() optimizer_d_aux.zero_grad() optimizer_d_main.zero_grad() # adapt LR if needed adjust_learning_rate(optimizer, i_iter, cfg) adjust_learning_rate_discriminator(optimizer_d_aux, i_iter, cfg) adjust_learning_rate_discriminator(optimizer_d_main, i_iter, cfg) # UDA Training # only train segnet. Don't accumulate grads in disciminators for param in d_aux.parameters(): param.requires_grad = False for param in d_main.parameters(): param.requires_grad = False # train on source _, batch = trainloader_iter.__next__() images_source, labels, _, _ = batch # debug: # labels=labels.numpy() # from matplotlib import pyplot as plt # import numpy as np # plt.figure(1), plt.imshow(labels[0]), plt.ion(), plt.colorbar(), plt.show() pred_src_aux, pred_src_main = model(images_source.cuda(device)) if cfg.TRAIN.MULTI_LEVEL: pred_src_aux = interp(pred_src_aux) loss_seg_src_aux = loss_calc(pred_src_aux, labels, device) else: loss_seg_src_aux = 0 pred_src_main = interp(pred_src_main) loss_seg_src_main = loss_calc(pred_src_main, labels, device) # pdb.set_trace() loss = (cfg.TRAIN.LAMBDA_SEG_MAIN * loss_seg_src_main + cfg.TRAIN.LAMBDA_SEG_AUX * loss_seg_src_aux) loss.backward() # adversarial training ot fool the discriminator _, batch = targetloader_iter.__next__() images, _, _, _ = batch pred_trg_aux, pred_trg_main = model(images.cuda(device)) if cfg.TRAIN.MULTI_LEVEL: pred_trg_aux = interp_target(pred_trg_aux) d_out_aux = d_aux(prob_2_entropy(F.softmax(pred_trg_aux))) loss_adv_trg_aux = bce_loss(d_out_aux, source_label) else: loss_adv_trg_aux = 0 pred_trg_main = interp_target(pred_trg_main) d_out_main = d_main(prob_2_entropy(F.softmax(pred_trg_main))) loss_adv_trg_main = bce_loss(d_out_main, source_label) loss = (cfg.TRAIN.LAMBDA_ADV_MAIN * loss_adv_trg_main + cfg.TRAIN.LAMBDA_ADV_AUX * loss_adv_trg_aux) loss = loss loss.backward() # Train discriminator networks # enable training mode on discriminator networks for param in d_aux.parameters(): param.requires_grad = True for param in d_main.parameters(): param.requires_grad = True # train with source if cfg.TRAIN.MULTI_LEVEL: pred_src_aux = pred_src_aux.detach() d_out_aux = d_aux(prob_2_entropy(F.softmax(pred_src_aux))) loss_d_aux = bce_loss(d_out_aux, source_label) loss_d_aux = loss_d_aux / 2 loss_d_aux.backward() pred_src_main = pred_src_main.detach() d_out_main = d_main(prob_2_entropy(F.softmax(pred_src_main))) loss_d_main = bce_loss(d_out_main, source_label) loss_d_main = loss_d_main / 2 loss_d_main.backward() # train with target if cfg.TRAIN.MULTI_LEVEL: pred_trg_aux = pred_trg_aux.detach() d_out_aux = d_aux(prob_2_entropy(F.softmax(pred_trg_aux))) loss_d_aux = bce_loss(d_out_aux, target_label) loss_d_aux = loss_d_aux / 2 loss_d_aux.backward() else: loss_d_aux = 0 pred_trg_main = pred_trg_main.detach() d_out_main = d_main(prob_2_entropy(F.softmax(pred_trg_main))) loss_d_main = bce_loss(d_out_main, target_label) loss_d_main = loss_d_main / 2 loss_d_main.backward() optimizer.step() if cfg.TRAIN.MULTI_LEVEL: optimizer_d_aux.step() optimizer_d_main.step() current_losses = {'loss_seg_src_aux': loss_seg_src_aux, 'loss_seg_src_main': loss_seg_src_main, 'loss_adv_trg_aux': loss_adv_trg_aux, 'loss_adv_trg_main': loss_adv_trg_main, 'loss_d_aux': loss_d_aux, 'loss_d_main': loss_d_main} print_losses(current_losses, i_iter) if i_iter % cfg.TRAIN.SAVE_PRED_EVERY == 0 and i_iter != 0: print('taking snapshot ...') print('exp =', cfg.TRAIN.SNAPSHOT_DIR) snapshot_dir = Path(cfg.TRAIN.SNAPSHOT_DIR) torch.save(model.state_dict(), snapshot_dir / f'model_{i_iter}.pth') torch.save(d_aux.state_dict(), snapshot_dir / f'model_{i_iter}_D_aux.pth') torch.save(d_main.state_dict(), snapshot_dir / f'model_{i_iter}_D_main.pth') if i_iter >= cfg.TRAIN.EARLY_STOP - 1: break sys.stdout.flush() # Visualize with tensorboard if viz_tensorboard: log_losses_tensorboard(writer, current_losses, i_iter) if i_iter % cfg.TRAIN.TENSORBOARD_VIZRATE == cfg.TRAIN.TENSORBOARD_VIZRATE - 1: draw_in_tensorboard(writer, images, i_iter, pred_trg_main, num_classes, 'T') draw_in_tensorboard(writer, images_source, i_iter, pred_src_main, num_classes, 'S')
def main(args): # load configuration file device = cfg.GPU_ID assert args.cfg is not None, 'Missing cfg file' cfg_from_file(args.cfg) if not os.path.exists(save_dir % (args.FDA_mode, args.round)): os.mkdir(save_dir % (args.FDA_mode, args.round)) # ----------------------------------------------------------------# args.LB = str(args.MBT) # set args.LB = 'MBT' SRC_IMG_MEAN = np.asarray(cfg.TRAIN.IMG_MEAN, dtype=np.float32) SRC_IMG_MEAN = torch.reshape(torch.from_numpy(SRC_IMG_MEAN), (1, 3, 1, 1)) ###################### here, replace by restoring three different model##################### if args.round == 0: # first round of SSL cfg.EXP_NAME = f'{cfg.SOURCE}2{cfg.TARGET}_{cfg.TRAIN.MODEL}_{cfg.TRAIN.DA_METHOD}_{args.FDA_mode}_LB_{args.LB}' elif args.round > 0: # when SSL round is higher than 0 # SOURCE and TARGET are no longer GTA and Cityscape, but are easy and hard split cfg.SOURCE = 'CityscapesEasy' cfg.TARGET = 'CityscapesHard' cfg.EXP_NAME = f'{cfg.SOURCE}2{cfg.TARGET}_{cfg.TRAIN.MODEL}_{cfg.TRAIN.DA_METHOD}_{args.FDA_mode}_LB_{args.LB}_THRESH_{str(thresholding)}_ROUND_{args.round - 1}' else: raise KeyError() ########################################################################################################################################## # ----------------------------------------------------------------# cfg.TEST.SNAPSHOT_DIR[0] = osp.join(cfg.EXP_ROOT_SNAPSHOT, cfg.EXP_NAME) # load model with parameters trained from Inter-domain adaptation model_gen = get_deeplab_v2(num_classes=cfg.NUM_CLASSES, multi_level=cfg.TEST.MULTI_LEVEL) restore_from = osp.join(cfg.TEST.SNAPSHOT_DIR[0], f'model_{args.best_iter}.pth') print("Loading the generator:", restore_from) load_checkpoint_for_evaluation(model_gen, restore_from, device) # load data target_dataset = CityscapesDataSet(args=args, root=cfg.DATA_DIRECTORY_TARGET, list_path=cfg.DATA_LIST_TARGET, set=cfg.TRAIN.SET_TARGET, info_path=cfg.TRAIN.INFO_TARGET, max_iters=None, crop_size=cfg.TRAIN.INPUT_SIZE_TARGET, mean=cfg.TRAIN.IMG_MEAN) target_loader = data.DataLoader(target_dataset, batch_size=cfg.TRAIN.BATCH_SIZE_TARGET, num_workers=cfg.NUM_WORKERS, shuffle=True, pin_memory=True, worker_init_fn=None) target_loader_iter = enumerate(target_loader) # upsampling layer input_size_target = cfg.TRAIN.INPUT_SIZE_TARGET interp_target = nn.Upsample(size=(input_size_target[1], input_size_target[0]), mode='bilinear', align_corners=True) # ---------------------------------------------------------------------------------------------------------------# # step 1. entropy-ranking: split the target dataset into easy and hard cases. entropy_list = [] for index in tqdm(range(len(target_loader))): _, batch = target_loader_iter.__next__() image, _, _, name = batch # normalize the image before fed into the trained model B, C, H, W = image.shape mean_image = SRC_IMG_MEAN.repeat(B, 1, H, W) if args.FDA_mode == 'on': image -= mean_image elif args.FDA_mode == 'off': # no need to perform normalization again since that has been done already in dataset class(GTA5, cityscapes) when args.FDA_mode = 'off' image = image else: raise KeyError() with torch.no_grad(): _, pred_trg_main = model_gen( image.cuda(device)) # shape(pred_trg_main) = (1, 19, 65, 129) pred_trg_main = interp_target( pred_trg_main) # shape(pred_trg_main) = (1, 19, 512, 1024) if args.normalize == True: normalizor = (11 - len(find_rare_class(pred_trg_main))) / 11.0 + 0.5 else: normalizor = 1 pred_trg_entropy = prob_2_entropy(F.softmax(pred_trg_main)) entropy_list.append( (name[0], pred_trg_entropy.mean().item() * normalizor)) # colorize_save(pred_trg_main, name[0], args.FDA_mode) # split the enntropy_list into _, easy_split = cluster_subdomain(entropy_list, args, thresholding) # ---------------------------------------------------------------------------------------------------------------# # step2. apply thresholding(either top 66% or confidence score above 0.9) to easy-split target dataset and save them. predicted_label = np.zeros( (len(easy_split), 512, 1024)) # (512, 1024) is the size of target output predicted_prob = np.zeros((len(easy_split), 512, 1024)) image_name = [] idx = 0 target_loader_iter = enumerate(target_loader) for index in tqdm(range(len(target_loader))): _, batch = target_loader_iter.__next__() image, _, _, name = batch if name[0] not in easy_split: # only compute the images that belongs to easy-split continue # normalize the image before fed into the trained model B, C, H, W = image.shape mean_image = SRC_IMG_MEAN.repeat(B, 1, H, W) if args.FDA_mode == 'on': image -= mean_image elif args.FDA_mode == 'off': # no need to perform normalization again since that has been done already in dataset class(GTA5, cityscapes) when args.FDA_mode = 'off' image = image else: raise KeyError() with torch.no_grad(): _, pred_trg_main = model_gen( image.cuda(device)) # shape(pred_trg_main) = (1, 19, 65, 129) pred_trg_main = F.softmax(interp_target(pred_trg_main), dim=1).cpu( ).data[0].numpy() # shape(pred_trg_main) = (1, 19, 512, 1024) pred_trg_main = pred_trg_main.transpose( 1, 2, 0) # shape(pred_trg_main) = (512, 1024, 19) label, prob = np.argmax(pred_trg_main, axis=2), np.max(pred_trg_main, axis=2) predicted_label[idx] = label predicted_prob[idx] = prob image_name.append(name[0]) idx += 1 assert len(easy_split) == len( image_name) # check whether all images in easy-split are processed # compute the threshold for each label thres = [] for i in range(cfg.NUM_CLASSES): x = predicted_prob[predicted_label == i] if len(x) == 0: thres.append(0) continue x = np.sort(x) thres.append( x[np.int(np.round(len(x) * 0.66))] ) # thres contains the thresholding values by labels in corresponding entry:thres[label] print(thres) thres = np.array(thres) thres[thres > 0.9] = 0.9 print(thres) colorize_save_with_thresholding(easy_split, thres, predicted_label, predicted_prob, image_name, args)
def main(args): # load configuration file device = cfg.GPU_ID assert args.cfg is not None, 'Missing cfg file' cfg_from_file(args.cfg) if not os.path.exists(save_dir % (args.FDA_mode)): os.mkdir(save_dir % (args.FDA_mode)) # ----------------------------------------------------------------# args.LB = str(args.LB).replace('.', '_') SRC_IMG_MEAN = np.asarray(cfg.TRAIN.IMG_MEAN, dtype=np.float32) SRC_IMG_MEAN = torch.reshape(torch.from_numpy(SRC_IMG_MEAN), (1, 3, 1, 1)) if args.round == 0: # first round of SSL cfg.EXP_NAME = f'{cfg.SOURCE}2{cfg.TARGET}_{cfg.TRAIN.MODEL}_{cfg.TRAIN.DA_METHOD}_{args.FDA_mode}_LB_{args.LB}' elif args.round > 0: # when SSL round is higher than 0 # SOURCE and TARGET are no longer GTA and Cityscape, but are easy and hard split cfg.SOURCE = 'CityscapesEasy' cfg.TARGET = 'CityscapesHard' cfg.EXP_NAME = f'{cfg.SOURCE}2{cfg.TARGET}_{cfg.TRAIN.MODEL}_{cfg.TRAIN.DA_METHOD}_{args.FDA_mode}_LB_{args.LB}_THRESH_{str(thresholding)}_ROUND_{args.round - 1}' else: raise KeyError() #cfg.EXP_NAME = f'{cfg.SOURCE}2{cfg.TARGET}_{cfg.TRAIN.MODEL}_{cfg.TRAIN.DA_METHOD}_{args.FDA_mode}_LB_{args.LB}' # ----------------------------------------------------------------# cfg.TEST.SNAPSHOT_DIR[0] = osp.join(cfg.EXP_ROOT_SNAPSHOT, cfg.EXP_NAME) # load model with parameters trained from Inter-domain adaptation model_gen = get_deeplab_v2(num_classes=cfg.NUM_CLASSES, multi_level=cfg.TEST.MULTI_LEVEL) restore_from = osp.join(cfg.TEST.SNAPSHOT_DIR[0], f'model_{args.best_iter}.pth') print("Loading the generator:", restore_from) load_checkpoint_for_evaluation(model_gen, restore_from, device) # load data target_dataset = CityscapesDataSet(args=args, root=cfg.DATA_DIRECTORY_TARGET, list_path=cfg.DATA_LIST_TARGET, set=cfg.TRAIN.SET_TARGET, info_path=cfg.TRAIN.INFO_TARGET, max_iters=None, crop_size=cfg.TRAIN.INPUT_SIZE_TARGET, mean=cfg.TRAIN.IMG_MEAN) target_loader = data.DataLoader(target_dataset, batch_size=cfg.TRAIN.BATCH_SIZE_TARGET, num_workers=cfg.NUM_WORKERS, shuffle=True, pin_memory=True, worker_init_fn=None) target_loader_iter = enumerate(target_loader) # upsampling layer input_size_target = cfg.TRAIN.INPUT_SIZE_TARGET interp_target = nn.Upsample(size=(input_size_target[1], input_size_target[0]), mode='bilinear', align_corners=True) entropy_list = [] for index in tqdm(range(len(target_loader))): _, batch = target_loader_iter.__next__() image, _, _, name = batch # ----------------------------------------------------------------# """ normalize the image before fed into the trained model """ B, C, H, W = image.shape mean_image = SRC_IMG_MEAN.repeat(B, 1, H, W) if args.FDA_mode == 'on': image -= mean_image elif args.FDA_mode == 'off': # no need to perform normalization again since that has been done already in dataset class(GTA5, cityscapes) when args.FDA_mode = 'off' image = image else: raise KeyError() # ----------------------------------------------------------------# with torch.no_grad(): _, pred_trg_main = model_gen(image.cuda(device)) pred_trg_main = interp_target(pred_trg_main) if args.normalize == True: normalizor = (11 - len(find_rare_class(pred_trg_main))) / 11.0 + 0.5 else: normalizor = 1 pred_trg_entropy = prob_2_entropy(F.softmax(pred_trg_main)) entropy_list.append( (name[0], pred_trg_entropy.mean().item() * normalizor)) colorize_save(pred_trg_main, name[0], args) # split the enntropy_list into cluster_subdomain(entropy_list, args, thresholding)
def main(args): # load configuration file device = cfg.GPU_ID assert args.cfg is not None, 'Missing cfg file' cfg_from_file(args.cfg) if not os.path.exists('./color_masks'): os.mkdir('./color_masks') cfg.EXP_NAME = f'{cfg.SOURCE}2{cfg.TARGET}_{cfg.TRAIN.MODEL}_{cfg.TRAIN.DA_METHOD}' cfg.TEST.SNAPSHOT_DIR[0] = osp.join(cfg.EXP_ROOT_SNAPSHOT, cfg.EXP_NAME) # load model with parameters trained from Inter-domain adaptation model_gen = get_deeplab_v2(num_classes=cfg.NUM_CLASSES, multi_level=cfg.TEST.MULTI_LEVEL) restore_from = osp.join(cfg.TEST.SNAPSHOT_DIR[0], f'model_{args.best_iter}.pth') print("Loading the generator:", restore_from) load_checkpoint_for_evaluation(model_gen, restore_from, device) # load data target_dataset = CityscapesDataSet(root=cfg.DATA_DIRECTORY_TARGET, list_path=cfg.DATA_LIST_TARGET, set=cfg.TRAIN.SET_TARGET, info_path=cfg.TRAIN.INFO_TARGET, max_iters=None, crop_size=cfg.TRAIN.INPUT_SIZE_TARGET, mean=cfg.TRAIN.IMG_MEAN) target_loader = data.DataLoader(target_dataset, batch_size=cfg.TRAIN.BATCH_SIZE_TARGET, num_workers=cfg.NUM_WORKERS, shuffle=True, pin_memory=True, worker_init_fn=None) target_loader_iter = enumerate(target_loader) # upsampling layer input_size_target = cfg.TRAIN.INPUT_SIZE_TARGET interp_target = nn.Upsample(size=(input_size_target[1], input_size_target[0]), mode='bilinear', align_corners=True) entropy_list = [] for index in tqdm(range(len(target_loader))): _, batch = target_loader_iter.__next__() image, _, _, name = batch with torch.no_grad(): _, pred_trg_main = model_gen(image.cuda(device)) pred_trg_main = interp_target(pred_trg_main) if args.normalize == True: normalizor = (11 - len(find_rare_class(pred_trg_main))) / 11.0 + 0.5 else: normalizor = 1 pred_trg_entropy = prob_2_entropy(F.softmax(pred_trg_main)) entropy_list.append( (name[0], pred_trg_entropy.mean().item() * normalizor)) colorize_save(pred_trg_main, name[0]) # split the enntropy_list into cluster_subdomain(entropy_list, args.lambda1)
def main(args): # load configuration file device = cfg.GPU_ID assert args.cfg is not None, 'Missing cfg file' cfg_from_file(args.cfg) if not os.path.exists('./output2'): os.mkdir('./output2') cfg.EXP_NAME = f'{cfg.SOURCE}2{cfg.TARGET}_{cfg.TRAIN.MODEL}_{cfg.TRAIN.DA_METHOD}_{cfg.NUM_CLASSES}class_18_01' cfg.TEST.SNAPSHOT_DIR[0] = osp.join(cfg.EXP_ROOT_SNAPSHOT, cfg.EXP_NAME) # load model with parameters trained from Inter-domain adaptation model_gen = get_deeplab_v2(num_classes=cfg.NUM_CLASSES, multi_level=cfg.TEST.MULTI_LEVEL) restore_from = osp.join(cfg.TEST.SNAPSHOT_DIR[0], f'model_{args.best_iter}.pth') #restore_from = r'C:\semseg\IntraDA\ADVENT\experiments\snapshots\SimRunway2RealRunway_DeepLabv2_AdvEnt_5class_v2\model_120000.pth' print("Loading the generator:", restore_from) load_checkpoint_for_evaluation(model_gen, args.checkpoint, device) # load data target_dataset = RealRunwayDataSet(args.data_root, list_path=args.list_path + '/{}.txt', set=cfg.TEST.SET_TARGET, info_path=cfg.TEST.INFO_TARGET, crop_size=cfg.TEST.INPUT_SIZE_TARGET, mean=cfg.TEST.IMG_MEAN, labels_size=cfg.TEST.OUTPUT_SIZE_TARGET) target_loader = data.DataLoader(target_dataset, batch_size=cfg.TRAIN.BATCH_SIZE_TARGET, num_workers=cfg.NUM_WORKERS, shuffle=True, pin_memory=True, worker_init_fn=None) target_loader_iter = enumerate(target_loader) # upsampling layer input_size_target = cfg.TRAIN.INPUT_SIZE_TARGET interp_target = nn.Upsample(size=(input_size_target[1], input_size_target[0]), mode='bilinear', align_corners=True) entropy_list = [] for index in tqdm(range(len(target_loader))): _, batch = target_loader_iter.__next__() image, _, _, name = batch with torch.no_grad(): _, pred_trg_main = model_gen(image.cuda(device)) pred_trg_main = interp_target(pred_trg_main) if args.normalize == True: normalizor = (11 - len(find_rare_class(pred_trg_main))) / 11.0 + 0.5 else: normalizor = 1 pred_trg_entropy = prob_2_entropy(F.softmax(pred_trg_main)) entropy_list.append( (name[0], pred_trg_entropy.mean().item() * normalizor)) colorize_save(pred_trg_main, name[0])
def train_preview(model, source_loader, target_loader, cfg, comet_exp): # UDA TRAINING ''' UDA training with advent ''' # Create the model and start the training. input_size_source = cfg.TRAIN.INPUT_SIZE_SOURCE input_size_target = cfg.TRAIN.INPUT_SIZE_TARGET device = cfg.GPU_ID num_classes = cfg.NUM_CLASSES # SEGMNETATION NETWORK model.train() model.to(device) cudnn.benchmark = True cudnn.enabled = True # DISCRIMINATOR NETWORK # feature-level d_aux = get_fc_discriminator(num_classes=num_classes) d_aux.train() d_aux.to(device) # seg maps, i.e. output, level d_main = get_fc_discriminator(num_classes=num_classes) d_main.train() d_main.to(device) # OPTIMIZERS # segnet's optimizer optimizer = optim.SGD(model.optim_parameters(cfg.TRAIN.LEARNING_RATE), lr=cfg.TRAIN.LEARNING_RATE, momentum=cfg.TRAIN.MOMENTUM, weight_decay=cfg.TRAIN.WEIGHT_DECAY) # discriminators' optimizers optimizer_d_aux = optim.Adam(d_aux.parameters(), lr=cfg.TRAIN.LEARNING_RATE_D, betas=(0.9, 0.99)) optimizer_d_main = optim.Adam(d_main.parameters(), lr=cfg.TRAIN.LEARNING_RATE_D, betas=(0.9, 0.99)) # interpolate output segmaps interp = nn.Upsample(size=(input_size_source[1], input_size_source[0]), mode='bilinear', align_corners=True) interp_target = nn.Upsample(size=(input_size_target[1], input_size_target[0]), mode='bilinear', align_corners=True) # labels for adversarial training source_label = 0 target_label = 1 times = deque([0], maxlen=100) model_times = deque([0], maxlen=100) source_loader_iter = enumerate(source_loader) target_loader_iter = enumerate(target_loader) cur_best_miou = -1 cur_best_model = '' for i_iter in tqdm(range(cfg.TRAIN.EARLY_STOP + 1)): times.append(time()) comet_exp.log_metric("i_iter", i_iter) comet_exp.log_metric("target_epoch", i_iter / len(target_loader)) comet_exp.log_metric("source_epoch", i_iter / len(source_loader)) # reset optimizers optimizer.zero_grad() optimizer_d_aux.zero_grad() optimizer_d_main.zero_grad() # adapt LR if needed adjust_learning_rate(optimizer, i_iter, cfg) adjust_learning_rate_discriminator(optimizer_d_aux, i_iter, cfg) adjust_learning_rate_discriminator(optimizer_d_main, i_iter, cfg) # UDA Training # only train segnet. Don't accumulate grads in disciminators for param in d_aux.parameters(): param.requires_grad = False for param in d_main.parameters(): param.requires_grad = False # train on source try: _, batch_and_path = source_loader_iter.__next__() except StopIteration: source_loader_iter = enumerate(source_loader) _, batch_and_path = source_loader_iter.__next__() images_source, labels = batch_and_path['data']['x'], batch_and_path[ 'data']['m'] pred_src_aux, pred_src_main = model(images_source.cuda(device)) if cfg.TRAIN.MULTI_LEVEL: pred_src_aux = interp(pred_src_aux) loss_seg_src_aux = loss_calc(pred_src_aux, labels, device) else: loss_seg_src_aux = 0 pred_src_main = interp(pred_src_main) loss_seg_src_main = loss_calc(pred_src_main, labels, device) loss = (cfg.TRAIN.LAMBDA_SEG_MAIN * loss_seg_src_main + cfg.TRAIN.LAMBDA_SEG_AUX * loss_seg_src_aux) loss.backward() # adversarial training to fool the discriminator try: _, batch = target_loader_iter.__next__() except StopIteration: target_loader_iter = enumerate(target_loader) _, batch = target_loader_iter.__next__() images = batch['data']['x'] pred_trg_aux, pred_trg_main = model(images.cuda(device)) if cfg.TRAIN.MULTI_LEVEL: pred_trg_aux = interp_target(pred_trg_aux) d_out_aux = d_aux(prob_2_entropy(F.softmax(pred_trg_aux))) loss_adv_trg_aux = bce_loss(d_out_aux, source_label) else: loss_adv_trg_aux = 0 pred_trg_main = interp_target(pred_trg_main) d_out_main = d_main(prob_2_entropy(F.softmax(pred_trg_main))) loss_adv_trg_main = bce_loss(d_out_main, source_label) loss = (cfg.TRAIN.LAMBDA_ADV_MAIN * loss_adv_trg_main + cfg.TRAIN.LAMBDA_ADV_AUX * loss_adv_trg_aux) loss = loss loss.backward() # Train discriminator networks # enable training mode on discriminator networks for param in d_aux.parameters(): param.requires_grad = True for param in d_main.parameters(): param.requires_grad = True # train with source if cfg.TRAIN.MULTI_LEVEL: pred_src_aux = pred_src_aux.detach() d_out_aux = d_aux(prob_2_entropy(F.softmax(pred_src_aux))) loss_d_aux = bce_loss(d_out_aux, source_label) loss_d_aux = loss_d_aux / 2 loss_d_aux.backward() pred_src_main = pred_src_main.detach() d_out_main = d_main(prob_2_entropy(F.softmax(pred_src_main))) loss_d_main = bce_loss(d_out_main, source_label) loss_d_main = loss_d_main / 2 loss_d_main.backward() # train with target if cfg.TRAIN.MULTI_LEVEL: pred_trg_aux = pred_trg_aux.detach() d_out_aux = d_aux(prob_2_entropy(F.softmax(pred_trg_aux))) loss_d_aux = bce_loss(d_out_aux, target_label) loss_d_aux = loss_d_aux / 2 loss_d_aux.backward() else: loss_d_aux = 0 pred_trg_main = pred_trg_main.detach() d_out_main = d_main(prob_2_entropy(F.softmax(pred_trg_main))) loss_d_main = bce_loss(d_out_main, target_label) loss_d_main = loss_d_main / 2 loss_d_main.backward() optimizer.step() if cfg.TRAIN.MULTI_LEVEL: optimizer_d_aux.step() optimizer_d_main.step() model_times.append(time() - times[-1]) mod_times = np.mean(model_times) comet_exp.log_metric("model_time", mod_times) current_losses = { 'loss_seg_src_aux': loss_seg_src_aux, 'loss_seg_src_main': loss_seg_src_main, 'loss_adv_trg_aux': loss_adv_trg_aux, 'loss_adv_trg_main': loss_adv_trg_main, 'loss_d_aux': loss_d_aux, 'loss_d_main': loss_d_main } print_losses(current_losses, i_iter) current_losses_numDict = tesnorDict2numDict(current_losses) comet_exp.log_metrics(current_losses_numDict) if i_iter % cfg.TRAIN.SAVE_PRED_EVERY == 0 and i_iter != 0: print('taking snapshot ...') print('exp =', cfg.TRAIN.SNAPSHOT_DIR) snapshot_dir = Path(cfg.TRAIN.SNAPSHOT_DIR) torch.save(model.state_dict(), snapshot_dir / f'model_{i_iter}.pth') torch.save(d_aux.state_dict(), snapshot_dir / f'model_{i_iter}_D_aux.pth') torch.save(d_main.state_dict(), snapshot_dir / f'model_{i_iter}_D_main.pth') if i_iter >= cfg.TRAIN.EARLY_STOP - 1: break if i_iter % cfg.TRAIN.SAVE_IMAGE_PRED == 0 and i_iter != 0 or i_iter == cfg.TRAIN.EARLY_STOP: print("Inferring test images in iteration {}...".format(i_iter)) hist = np.zeros((cfg.NUM_CLASSES, cfg.NUM_CLASSES)) image, label = batch['data']['x'][0], batch['data']['m'][0] image = image[None, :, :, :] interp = nn.Upsample(size=(label.shape[1], label.shape[2]), mode='bilinear', align_corners=True) with torch.no_grad(): pred_main = model(image.cuda(device))[1] output = interp(pred_main).cpu().data[0].numpy() output = output.transpose(1, 2, 0) output = np.argmax(output, axis=2) label0 = label.numpy()[0] hist += fast_hist(label0.flatten(), output.flatten(), cfg.NUM_CLASSES) output = torch.tensor(output, dtype=torch.float32) output = output[None, :, :] output_RGB = output.repeat(3, 1, 1) if i_iter % 100 == 0: print('{:d} / {:d}: {:0.2f}'.format( i_iter % len(target_loader), len(target_loader), 100 * np.nanmean(per_class_iu(hist)))) inters_over_union_classes = per_class_iu(hist) computed_miou = round( np.nanmean(inters_over_union_classes) * 100, 2) if cur_best_miou < computed_miou: cur_best_miou = computed_miou cur_best_model = f'model_{i_iter}.pth' print('\tCurrent mIoU:', computed_miou) print('\tCurrent best model:', cur_best_model) print('\tCurrent best mIoU:', cur_best_miou) mious = { 'Current mIoU': computed_miou, 'Current best model': cur_best_model, 'Current best mIoU': cur_best_miou } comet_exp.log_metrics(mious) image = image[0] # change size from [1,x,y,z] to [x,y,z] save_images = [] save_images.append(image) # Overlay mask: save_mask = (image - (image * label.repeat(3, 1, 1)) + label.repeat(3, 1, 1)) save_fake_mask = (image - (image * output_RGB) + output_RGB) save_images.append(save_mask) save_images.append(save_fake_mask) save_images.append(label.repeat(3, 1, 1)) save_images.append(output_RGB) write_images(save_images, i_iter, comet_exp=comet_exp, store_im=cfg.TEST.store_images)
def train_dada(model, trainloader, targetloader, cfg): """ UDA training with dada """ # Create the model and start the training. input_size_source = cfg.TRAIN.INPUT_SIZE_SOURCE input_size_target = cfg.TRAIN.INPUT_SIZE_TARGET device = cfg.GPU_ID num_classes = cfg.NUM_CLASSES viz_tensorboard = os.path.exists(cfg.TRAIN.TENSORBOARD_LOGDIR) if viz_tensorboard: writer = SummaryWriter(log_dir=cfg.TRAIN.TENSORBOARD_LOGDIR) # SEGMNETATION NETWORK model.train() model.to(device) cudnn.benchmark = True cudnn.enabled = True # DISCRIMINATOR NETWORK # seg maps, i.e. output, level d_main = get_fc_discriminator(num_classes=num_classes) d_main.train() d_main.to(device) # OPTIMIZERS # segnet's optimizer optimizer = optim.SGD( model.optim_parameters(cfg.TRAIN.LEARNING_RATE), lr=cfg.TRAIN.LEARNING_RATE, momentum=cfg.TRAIN.MOMENTUM, weight_decay=cfg.TRAIN.WEIGHT_DECAY, ) # discriminators' optimizers optimizer_d_main = optim.Adam(d_main.parameters(), lr=cfg.TRAIN.LEARNING_RATE_D, betas=(0.9, 0.99)) # interpolate output segmaps interp = nn.Upsample( size=(input_size_source[1], input_size_source[0]), mode="bilinear", align_corners=True, ) interp_target = nn.Upsample( size=(input_size_target[1], input_size_target[0]), mode="bilinear", align_corners=True, ) # labels for adversarial training source_label = 0 target_label = 1 trainloader_iter = enumerate(trainloader) targetloader_iter = enumerate(targetloader) for i_iter in tqdm(range(cfg.TRAIN.EARLY_STOP + 1)): # reset optimizers optimizer.zero_grad() optimizer_d_main.zero_grad() adjust_learning_rate(optimizer, i_iter, cfg) adjust_learning_rate_discriminator(optimizer_d_main, i_iter, cfg) # UDA Training # only train segnet. Don't accumulate grads in disciminators for param in d_main.parameters(): param.requires_grad = False # train on source _, batch = trainloader_iter.__next__() images_source, labels, depth, _, _ = batch _, pred_src_main, pred_depth_src_main = model( images_source.cuda(device)) pred_src_main = interp(pred_src_main) pred_depth_src_main = interp(pred_depth_src_main) loss_depth_src_main = loss_calc_depth(pred_depth_src_main, depth, device) loss_seg_src_main = loss_calc(pred_src_main, labels, device) loss = (cfg.TRAIN.LAMBDA_SEG_MAIN * loss_seg_src_main + cfg.TRAIN.LAMBDA_DEPTH_MAIN * loss_depth_src_main) loss.backward() # adversarial training ot fool the discriminator _, batch = targetloader_iter.__next__() images, _, _, _ = batch _, pred_trg_main, pred_depth_trg_main = model(images.cuda(device)) pred_trg_main = interp_target(pred_trg_main) pred_depth_trg_main = interp_target(pred_depth_trg_main) d_out_main = d_main( prob_2_entropy(F.softmax(pred_trg_main)) * pred_depth_trg_main) loss_adv_trg_main = bce_loss(d_out_main, source_label) loss = cfg.TRAIN.LAMBDA_ADV_MAIN * loss_adv_trg_main loss.backward() # Train discriminator networks # enable training mode on discriminator networks for param in d_main.parameters(): param.requires_grad = True # train with source pred_src_main = pred_src_main.detach() pred_depth_src_main = pred_depth_src_main.detach() d_out_main = d_main( prob_2_entropy(F.softmax(pred_src_main)) * pred_depth_src_main) loss_d_main = bce_loss(d_out_main, source_label) loss_d_main = loss_d_main loss_d_main.backward() # train with target pred_trg_main = pred_trg_main.detach() pred_depth_trg_main = pred_depth_trg_main.detach() d_out_main = d_main( prob_2_entropy(F.softmax(pred_trg_main)) * pred_depth_trg_main) loss_d_main = bce_loss(d_out_main, target_label) loss_d_main = loss_d_main loss_d_main.backward() optimizer.step() optimizer_d_main.step() current_losses = { "loss_seg_src_main": loss_seg_src_main, "loss_depth_src_main": loss_depth_src_main, "loss_adv_trg_main": loss_adv_trg_main, "loss_d_main": loss_d_main, } print_losses(current_losses, i_iter) if i_iter % cfg.TRAIN.SAVE_PRED_EVERY == 0 and i_iter != 0: print("taking snapshot ...") print("exp =", cfg.TRAIN.SNAPSHOT_DIR) snapshot_dir = Path(cfg.TRAIN.SNAPSHOT_DIR) torch.save(model.state_dict(), snapshot_dir / f"model_{i_iter}.pth") torch.save(d_main.state_dict(), snapshot_dir / f"model_{i_iter}_D_main.pth") if i_iter >= cfg.TRAIN.EARLY_STOP - 1: break sys.stdout.flush() # Visualize with tensorboard if viz_tensorboard: log_losses_tensorboard(writer, current_losses, i_iter) if i_iter % cfg.TRAIN.TENSORBOARD_VIZRATE == cfg.TRAIN.TENSORBOARD_VIZRATE - 1: draw_in_tensorboard(writer, images, i_iter, pred_trg_main, num_classes, "T") draw_in_tensorboard(writer, images_source, i_iter, pred_src_main, num_classes, "S")
def train_advent(model, trainloader, targetloader, cfg, args): ''' UDA training with advent ''' # Create the model and start the training. # pdb.set_trace() input_size_source = cfg.TRAIN.INPUT_SIZE_SOURCE input_size_target = cfg.TRAIN.INPUT_SIZE_TARGET SRC_IMG_MEAN = np.asarray(cfg.TRAIN.IMG_MEAN, dtype=np.float32) SRC_IMG_MEAN = torch.reshape(torch.from_numpy(SRC_IMG_MEAN), (1, 3, 1, 1)) device = cfg.GPU_ID num_classes = cfg.NUM_CLASSES viz_tensorboard = os.path.exists(cfg.TRAIN.TENSORBOARD_LOGDIR) if viz_tensorboard: writer = SummaryWriter(log_dir=cfg.TRAIN.TENSORBOARD_LOGDIR) # -------------------------------------------------------- # # codes to initialize wandb for storing logs on its cloud wandb.init(project='FDA_integration_to_INTRA_DA') wandb.config.update(args) for key, val in cfg.items(): wandb.config.update({key: val}) wandb.watch(model) # -------------------------------------------------------- # # SEGMNETATION NETWORK model.train() model.to(device) cudnn.benchmark = True cudnn.enabled = True # DISCRIMINATOR NETWORK # feature-level d_aux = get_fc_discriminator(num_classes=num_classes) d_aux.train() d_aux.to(device) # restore_from = cfg.TRAIN.RESTORE_FROM_aux # print("Load Discriminator:", restore_from) # load_checkpoint_for_evaluation(d_aux, restore_from, device) # seg maps, i.e. output, level d_main = get_fc_discriminator(num_classes=num_classes) d_main.train() d_main.to(device) # restore_from = cfg.TRAIN.RESTORE_FROM_main # print("Load Discriminator:", restore_from) # load_checkpoint_for_evaluation(d_main, restore_from, device) # OPTIMIZERS # segnet's optimizer optimizer = optim.SGD(model.optim_parameters(cfg.TRAIN.LEARNING_RATE), lr=cfg.TRAIN.LEARNING_RATE, momentum=cfg.TRAIN.MOMENTUM, weight_decay=cfg.TRAIN.WEIGHT_DECAY) # discriminators' optimizers optimizer_d_aux = optim.Adam(d_aux.parameters(), lr=cfg.TRAIN.LEARNING_RATE_D, betas=(0.9, 0.99)) optimizer_d_main = optim.Adam(d_main.parameters(), lr=cfg.TRAIN.LEARNING_RATE_D, betas=(0.9, 0.99)) # interpolate output segmaps interp = nn.Upsample(size=(input_size_source[1], input_size_source[0]), mode='bilinear', align_corners=True) interp_target = nn.Upsample(size=(input_size_target[1], input_size_target[0]), mode='bilinear', align_corners=True) # labels for adversarial training source_label = 0 target_label = 1 trainloader_iter = enumerate(trainloader) targetloader_iter = enumerate(targetloader) for i_iter in tqdm(range(cfg.TRAIN.EARLY_STOP + 1)): # reset optimizers optimizer.zero_grad() optimizer_d_aux.zero_grad() optimizer_d_main.zero_grad() # adapt LR if needed adjust_learning_rate(optimizer, i_iter, cfg) adjust_learning_rate_discriminator(optimizer_d_aux, i_iter, cfg) adjust_learning_rate_discriminator(optimizer_d_main, i_iter, cfg) # UDA Training # only train segnet. Don't accumulate grads in disciminators for param in d_aux.parameters(): param.requires_grad = False for param in d_main.parameters(): param.requires_grad = False _, batch = trainloader_iter.__next__() images_source, labels, _, _ = batch _, batch = targetloader_iter.__next__() images, _, _, _ = batch # ----------------------------------------------------------------# B, C, H, W = images_source.shape mean_images_source = SRC_IMG_MEAN.repeat(B, 1, H, W) mean_images = SRC_IMG_MEAN.repeat(B, 1, H, W) if args.FDA_mode == 'on': # normalize the source and target image images_source -= mean_images_source images -= mean_images elif args.FDA_mode == 'off': # Keep source and target images as they are # no need to perform normalization again since that has been done already in dataset class(GTA5, cityscapes) when args.FDA_mode = 'off' images_source = images_source images = images else: raise KeyError() # ----------------------------------------------------------------# # debug: # labels=labels.numpy() # from matplotlib import pyplot as plt # import numpy as np # plt.figure(1), plt.imshow(labels[0]), plt.ion(), plt.colorbar(), plt.show() # train on source pred_src_aux, pred_src_main = model(images_source.cuda(device)) if cfg.TRAIN.MULTI_LEVEL: pred_src_aux = interp(pred_src_aux) loss_seg_src_aux = loss_calc(pred_src_aux, labels, device) else: loss_seg_src_aux = 0 pred_src_main = interp(pred_src_main) loss_seg_src_main = loss_calc(pred_src_main, labels, device) # pdb.set_trace() loss = (cfg.TRAIN.LAMBDA_SEG_MAIN * loss_seg_src_main + cfg.TRAIN.LAMBDA_SEG_AUX * loss_seg_src_aux) loss.backward() # adversarial training ot fool the discriminator pred_trg_aux, pred_trg_main = model(images.cuda(device)) if cfg.TRAIN.MULTI_LEVEL: pred_trg_aux = interp_target(pred_trg_aux) d_out_aux = d_aux(prob_2_entropy(F.softmax(pred_trg_aux))) loss_adv_trg_aux = bce_loss(d_out_aux, source_label) else: loss_adv_trg_aux = 0 pred_trg_main = interp_target(pred_trg_main) d_out_main = d_main(prob_2_entropy(F.softmax(pred_trg_main))) loss_adv_trg_main = bce_loss(d_out_main, source_label) loss = (cfg.TRAIN.LAMBDA_ADV_MAIN * loss_adv_trg_main + cfg.TRAIN.LAMBDA_ADV_AUX * loss_adv_trg_aux) loss = loss loss.backward() # Train discriminator networks # enable training mode on discriminator networks for param in d_aux.parameters(): param.requires_grad = True for param in d_main.parameters(): param.requires_grad = True # train with source if cfg.TRAIN.MULTI_LEVEL: pred_src_aux = pred_src_aux.detach() d_out_aux = d_aux(prob_2_entropy(F.softmax(pred_src_aux))) loss_d_aux = bce_loss(d_out_aux, source_label) loss_d_aux = loss_d_aux / 2 loss_d_aux.backward() pred_src_main = pred_src_main.detach() d_out_main = d_main(prob_2_entropy(F.softmax(pred_src_main))) loss_d_main = bce_loss(d_out_main, source_label) loss_d_main = loss_d_main / 2 loss_d_main.backward() # train with target if cfg.TRAIN.MULTI_LEVEL: pred_trg_aux = pred_trg_aux.detach() d_out_aux = d_aux(prob_2_entropy(F.softmax(pred_trg_aux))) loss_d_aux = bce_loss(d_out_aux, target_label) loss_d_aux = loss_d_aux / 2 loss_d_aux.backward() else: loss_d_aux = 0 pred_trg_main = pred_trg_main.detach() d_out_main = d_main(prob_2_entropy(F.softmax(pred_trg_main))) loss_d_main = bce_loss(d_out_main, target_label) loss_d_main = loss_d_main / 2 loss_d_main.backward() optimizer.step() if cfg.TRAIN.MULTI_LEVEL: optimizer_d_aux.step() optimizer_d_main.step() if i_iter % cfg.TRAIN.SAVE_PRED_EVERY == 0 and i_iter != 0: print('taking snapshot ...') print('exp =', cfg.TRAIN.SNAPSHOT_DIR) snapshot_dir = Path(cfg.TRAIN.SNAPSHOT_DIR) torch.save(model.state_dict(), snapshot_dir / f'model_{i_iter}.pth') torch.save(d_aux.state_dict(), snapshot_dir / f'model_{i_iter}_D_aux.pth') torch.save(d_main.state_dict(), snapshot_dir / f'model_{i_iter}_D_main.pth') if i_iter >= cfg.TRAIN.EARLY_STOP - 1: break sys.stdout.flush() # Visualize with tensorboard if viz_tensorboard: # ----------------------------------------------------------------# if i_iter % cfg.TRAIN.TENSORBOARD_VIZRATE == cfg.TRAIN.TENSORBOARD_VIZRATE - 1: current_losses = { 'loss_seg_src_aux': loss_seg_src_aux, 'loss_seg_src_main': loss_seg_src_main, 'loss_adv_trg_aux': loss_adv_trg_aux, 'loss_adv_trg_main': loss_adv_trg_main, 'loss_d_aux': loss_d_aux, 'loss_d_main': loss_d_main } print_losses(current_losses, i_iter) log_losses_tensorboard(writer, current_losses, i_iter) draw_in_tensorboard(writer, images + mean_images, i_iter, pred_trg_main, num_classes, 'T') draw_in_tensorboard(writer, images_source + mean_images_source, i_iter, pred_src_main, num_classes, 'S') wandb.log({'loss': current_losses}, step=(i_iter + 1)) if i_iter % (cfg.TRAIN.TENSORBOARD_VIZRATE == cfg.TRAIN.TENSORBOARD_VIZRATE ) * 25 - 1: # for every 2500 iteration wandb.log( {'source': wandb.Image(torch.flip(images_source+mean_images_source, [1]).cpu().data[0].numpy().transpose((1, 2, 0))), \ 'target': wandb.Image(torch.flip(images+mean_images, [1]).cpu().data[0].numpy().transpose((1, 2, 0))), 'pesudo label': wandb.Image(np.asarray(colorize_mask(np.asarray(labels.cpu().data.numpy().transpose(1,2,0).reshape((512,1024)), dtype=np.uint8)).convert('RGB')) )}, step=(i_iter + 1))