def iterate(mode, args, loader, model, optimizer, logger, epoch): block_average_meter = AverageMeter() average_meter = AverageMeter() meters = [block_average_meter, average_meter] # switch to appropriate mode assert mode in ["train", "val", "eval", "test_prediction", "test_completion"], \ "unsupported mode: {}".format(mode) if mode == 'train': model.train() lr = helper.adjust_learning_rate(args.lr, optimizer, epoch) else: model.eval() lr = 0 for i, batch_data in enumerate(loader): start = time.time() batch_data = { key: val.to(device) for key, val in batch_data.items() if val is not None } gt = batch_data[ 'gt'] if mode != 'test_prediction' and mode != 'test_completion' else None data_time = time.time() - start start = time.time() with torch.no_grad(): # 自己加的 pred = model(batch_data) depth_loss, photometric_loss, smooth_loss, mask = 0, 0, 0, None gpu_time = time.time() - start # measure accuracy and record loss with torch.no_grad(): mini_batch_size = next(iter(batch_data.values())).size(0) result = Result() if mode != 'test_prediction' and mode != 'test_completion': #result.evaluate(pred.data, gt.data, photometric_loss) result.evaluate(pred.data.cpu(), gt.data.cpu(), photometric_loss) [ m.update(result, gpu_time, data_time, mini_batch_size) for m in meters ] logger.conditional_print(mode, i, epoch, lr, len(loader), block_average_meter, average_meter) logger.conditional_save_img_comparison(mode, i, batch_data, pred, epoch) logger.conditional_save_pred(mode, i, pred, epoch) avg = logger.conditional_save_info(mode, average_meter, epoch) is_best = logger.rank_conditional_save_best(mode, avg, epoch) if is_best and not (mode == "train"): logger.save_img_comparison_as_best(mode, epoch) logger.conditional_summarize(mode, avg, is_best) return avg, is_best
def iterate(mode, args, loader, model, optimizer, logger, epoch): block_average_meter = AverageMeter() average_meter = AverageMeter() meters = [block_average_meter, average_meter] # switch to appropriate mode assert mode in ["train", "val", "eval", "test_prediction", "test_completion"], \ "unsupported mode: {}".format(mode) if mode == 'train': model.train() lr = helper.adjust_learning_rate(args.lr, optimizer, epoch) else: model.eval( ) # batchnorm or dropout layers will work in eval mode instead of training mode lr = 0 for i, batch_data in enumerate( loader ): # batch_data keys: 'd' (depth), 'gt' (ground truth), 'g' (gray) start = time.time() batch_data = { key: val.to(device) for key, val in batch_data.items() if val is not None } gt = batch_data[ 'gt'] if mode != 'test_prediction' and mode != 'test_completion' else None data_time = time.time() - start start = time.time() pred = model(batch_data) if args.save_images: # save depth predictions pred_out_dir = max(glob.glob('../outputs/var_final_NN/var.test*'), key=os.path.getmtime) + '/dense_depth_images' pred1 = pred.cpu().detach().numpy()[:, 0, :, :] for im_idx, pred_im in enumerate(pred1): pred_out_dir1 = os.path.abspath(pred_out_dir) cur_path = os.path.abspath((loader.dataset.paths)['d'][i]) basename = os.path.basename(cur_path) cur_dir = os.path.abspath(os.path.dirname(cur_path)) cur_dir = cur_dir.split('var_final_NN/')[1] new_dir = os.path.abspath(pred_out_dir1 + '/' + cur_dir) new_path = os.path.abspath(new_dir + '/' + basename) if os.path.isdir(new_dir) == False: os.makedirs(new_dir) depth_write(new_path, pred_im) depth_loss, photometric_loss, smooth_loss, mask = 0, 0, 0, None if mode == 'train': # Loss 1: the direct depth supervision from ground truth label # mask=1 indicates that a pixel does not ground truth labels if 'sparse' in args.train_mode: depth_loss = depth_criterion(pred, batch_data['d']) mask = (batch_data['d'] < 1e-3).float() elif 'dense' in args.train_mode: depth_loss = depth_criterion(pred, gt) mask = (gt < 1e-3).float() # Loss 2: the self-supervised photometric loss if args.use_pose: # create multi-scale pyramids pred_array = helper.multiscale(pred) rgb_curr_array = helper.multiscale(batch_data['rgb']) rgb_near_array = helper.multiscale(batch_data['rgb_near']) if mask is not None: mask_array = helper.multiscale(mask) num_scales = len(pred_array) # compute photometric loss at multiple scales for scale in range(len(pred_array)): pred_ = pred_array[scale] rgb_curr_ = rgb_curr_array[scale] rgb_near_ = rgb_near_array[scale] mask_ = None if mask is not None: mask_ = mask_array[scale] # compute the corresponding intrinsic parameters height_, width_ = pred_.size(2), pred_.size(3) intrinsics_ = kitti_intrinsics.scale(height_, width_) # inverse warp from a nearby frame to the current frame warped_ = homography_from(rgb_near_, pred_, batch_data['r_mat'], batch_data['t_vec'], intrinsics_) photometric_loss += photometric_criterion( rgb_curr_, warped_, mask_) * (2**(scale - num_scales)) # Loss 3: the depth smoothness loss smooth_loss = smoothness_criterion(pred) if args.w2 > 0 else 0 # backprop loss = depth_loss + args.w1 * photometric_loss + args.w2 * smooth_loss optimizer.zero_grad() loss.backward() optimizer.step() gpu_time = time.time() - start # measure accuracy and record loss of each batch with torch.no_grad( ): # impacts the autograd engine and deactivate it (will reduce memory usage and speed up computations) mini_batch_size = next(iter(batch_data.values())).size(0) result = Result() # metrics if mode != 'test_prediction' and mode != 'test_completion': result.evaluate(pred.data, gt.data, photometric_loss) [ m.update(result, gpu_time, data_time, mini_batch_size) for m in meters ] logger.conditional_print(mode, i, epoch, args.epochs, lr, len(loader), block_average_meter, average_meter) logger.conditional_save_img_comparison(mode, i, batch_data, pred, epoch) logger.conditional_save_pred(mode, i, pred, epoch) del pred avg = logger.conditional_save_info( mode, average_meter, epoch) # take the avg of all the batches, to get the epoch metrics is_best = logger.rank_conditional_save_best(mode, avg, epoch, args.epochs) if is_best and not (mode == "train"): logger.save_img_comparison_as_best(mode, epoch) logger.conditional_summarize(mode, avg, is_best) return avg, is_best
def iterate(mode, args, loader, model, optimizer, logger, epoch): block_average_meter = AverageMeter() average_meter = AverageMeter() meters = [block_average_meter, average_meter] # switch to appropriate mode assert mode in ["train", "val", "eval", "test_prediction", "test_completion"], \ "unsupported mode: {}".format(mode) if mode == 'train': model.train() lr = helper.adjust_learning_rate(args.lr, optimizer, epoch) else: model.eval() lr = 0 for i, batch_data in enumerate(loader): print("The batch data keys are {}".format(batch_data.keys())) start = time.time() batch_data = { key: val.to(device) for key, val in batch_data.items() if val is not None } gt = batch_data[ 'gt'] if mode != 'test_prediction' and mode != 'test_completion' else None data_time = time.time() - start start = time.time() temp_d = batch_data['d'] temp_gt = batch_data['gt'] temp_g = batch_data['g'] print("The depth min:{}, max:{}, shape:{}, dtype:{}".format( torch.min(temp_d), torch.max(temp_d), temp_d.shape, temp_d.dtype)) print("The groundtruth min:{}, max:{}, shape:{}, dtype:{}".format( torch.min(temp_gt), torch.max(temp_gt), temp_gt.shape, temp_gt.dtype)) print("The greyscale min:{}, max:{}, shape:{}, dtype:{}".format( torch.min(temp_g), torch.max(temp_g), temp_g.shape, temp_g.dtype)) pred = model(batch_data) temp_out = pred.detach().cpu() print("The output min:{}, max:{}, shape:{}, dtype:{}".format( torch.min(temp_out), torch.max(temp_out), temp_out.shape, temp_out.dtype)) depth_loss, photometric_loss, smooth_loss, mask = 0, 0, 0, None if mode == 'train': # Loss 1: the direct depth supervision from ground truth label # mask=1 indicates that a pixel does not ground truth labels if 'sparse' in args.train_mode: depth_loss = depth_criterion(pred, batch_data['d']) mask = (batch_data['d'] < 1e-3).float() elif 'dense' in args.train_mode: depth_loss = depth_criterion(pred, gt) mask = (gt < 1e-3).float() # Loss 2: the self-supervised photometric loss if args.use_pose: # create multi-scale pyramids pred_array = helper.multiscale(pred) rgb_curr_array = helper.multiscale(batch_data['rgb']) rgb_near_array = helper.multiscale(batch_data['rgb_near']) if mask is not None: mask_array = helper.multiscale(mask) num_scales = len(pred_array) # compute photometric loss at multiple scales for scale in range(len(pred_array)): pred_ = pred_array[scale] rgb_curr_ = rgb_curr_array[scale] rgb_near_ = rgb_near_array[scale] mask_ = None if mask is not None: mask_ = mask_array[scale] # compute the corresponding intrinsic parameters height_, width_ = pred_.size(2), pred_.size(3) intrinsics_ = kitti_intrinsics.scale(height_, width_) # inverse warp from a nearby frame to the current frame warped_ = homography_from(rgb_near_, pred_, batch_data['r_mat'], batch_data['t_vec'], intrinsics_) photometric_loss += photometric_criterion( rgb_curr_, warped_, mask_) * (2**(scale - num_scales)) # Loss 3: the depth smoothness loss smooth_loss = smoothness_criterion(pred) if args.w2 > 0 else 0 # backprop loss = depth_loss + args.w1 * photometric_loss + args.w2 * smooth_loss optimizer.zero_grad() loss.backward() optimizer.step() gpu_time = time.time() - start # measure accuracy and record loss with torch.no_grad(): mini_batch_size = next(iter(batch_data.values())).size(0) result = Result() if mode != 'test_prediction' and mode != 'test_completion': result.evaluate(pred.data, gt.data, photometric_loss) [ m.update(result, gpu_time, data_time, mini_batch_size) for m in meters ] logger.conditional_print(mode, i, epoch, lr, len(loader), block_average_meter, average_meter) logger.conditional_save_img_comparison(mode, i, batch_data, pred, epoch) logger.conditional_save_pred(mode, i, pred, epoch) avg = logger.conditional_save_info(mode, average_meter, epoch) is_best = logger.rank_conditional_save_best(mode, avg, epoch) if is_best and not (mode == "train"): logger.save_img_comparison_as_best(mode, epoch) logger.conditional_summarize(mode, avg, is_best) return avg, is_best
def iterate(mode, args, loader, model, optimizer, logger, epoch): block_average_meter = AverageMeter() average_meter = AverageMeter() meters = [block_average_meter, average_meter] # switch to appropriate mode assert mode in ["train", "val", "eval", "test_prediction", "test_completion"], \ "unsupported mode: {}".format(mode) if mode == 'train': model.train() lr = helper.adjust_learning_rate(args.lr, optimizer, epoch) else: model.eval() lr = 0 torch.set_printoptions(profile="full") table_is = np.zeros(400) for i, batch_data in enumerate(loader): sparse_depth_pathname = batch_data['d_path'][0] print(sparse_depth_pathname) del batch_data['d_path'] print("i: ", i) start = time.time() batch_data = { key: val.to(device) for key, val in batch_data.items() if val is not None } gt = batch_data[ 'gt'] if mode != 'test_prediction' and mode != 'test_completion' else None # adjust depth for features depth_adjust = args.depth_adjust adjust_features = False if depth_adjust and args.use_d: if args.type_feature == "sq": if args.use_rgb: depth_new, alg_mode, feat_mode, features, shape = depth_adjustment( batch_data['d'], args.test_mode, args.feature_mode, args.feature_num, args.rank_file_global_sq, adjust_features, i, model_orig, args.seed, batch_data['rgb']) else: depth_new, alg_mode, feat_mode, features, shape = depth_adjustment( batch_data['d'], args.test_mode, args.feature_mode, args.feature_num, args.rank_file_global_sq, adjust_features, i, model_orig, args.seed) elif args.type_feature == "lines": depth_new, alg_mode, feat_mode, features = depth_adjustment_lines( batch_data['d'], args.test_mode, args.feature_mode, args.feature_num, args.rank_file_global_sq, i, model_orig, args.seed) batch_data['d'] = torch.Tensor(depth_new).unsqueeze(0).unsqueeze( 1).to(device) data_time = time.time() - start start = time.time() if mode == "train": pred = model(batch_data) else: with torch.no_grad(): pred = model(batch_data) # im = batch_data['d'].detach().cpu().numpy() # im_sq = im.squeeze() # plt.figure() # plt.imshow(im_sq) # plt.show() # for i in range(im_sq.shape[0]): # print(f"{i} - {np.sum(im_sq[i])}") # pred = pred +0.155 # gt = gt+0.155 # compute loss depth_loss, photometric_loss, smooth_loss, mask = 0, 0, 0, None if mode == 'train': # Loss 1: the direct depth supervision from ground truth label # mask=1 indicates that a pixel does not ground truth labels if 'sparse' in args.train_mode: depth_loss = depth_criterion(pred, batch_data['d']) mask = (batch_data['d'] < 1e-3).float() elif 'dense' in args.train_mode: depth_loss = depth_criterion(pred, gt) mask = (gt < 1e-3).float() # Loss 2: the self-supervised photometric loss if args.use_pose: # create multi-scale pyramids pred_array = helper.multiscale(pred) rgb_curr_array = helper.multiscale(batch_data['rgb']) rgb_near_array = helper.multiscale(batch_data['rgb_near']) if mask is not None: mask_array = helper.multiscale(mask) num_scales = len(pred_array) # compute photometric loss at multiple scales for scale in range(len(pred_array)): pred_ = pred_array[scale] rgb_curr_ = rgb_curr_array[scale] rgb_near_ = rgb_near_array[scale] mask_ = None if mask is not None: mask_ = mask_array[scale] # compute the corresponding intrinsic parameters height_, width_ = pred_.size(2), pred_.size(3) intrinsics_ = kitti_intrinsics.scale(height_, width_) # inverse warp from a nearby frame to the current frame warped_ = homography_from(rgb_near_, pred_, batch_data['r_mat'], batch_data['t_vec'], intrinsics_) photometric_loss += photometric_criterion( rgb_curr_, warped_, mask_) * (2**(scale - num_scales)) # Loss 3: the depth smoothness loss smooth_loss = smoothness_criterion(pred) if args.w2 > 0 else 0 # backprop loss = depth_loss + args.w1 * photometric_loss + args.w2 * smooth_loss optimizer.zero_grad() loss.backward() optimizer.step() gpu_time = time.time() - start # measure accuracy and record loss with torch.no_grad(): mini_batch_size = next(iter(batch_data.values())).size(0) result = Result() if mode != 'test_prediction' and mode != 'test_completion': result.evaluate(pred.data, gt.data, photometric_loss) [ m.update(result, gpu_time, data_time, mini_batch_size) for m in meters ] print(f"rmse: {result.rmse:,}") if result.rmse < 6000: print("good rmse") elif result.rmse > 12000: print("bad rmse") logger.conditional_print(mode, i, epoch, lr, len(loader), block_average_meter, average_meter) logger.conditional_save_img_comparison(mode, i, batch_data, pred, epoch) logger.conditional_save_pred(mode, i, pred, epoch) # save log and checkpoint every = 999 if mode == "val" else 200 if i % every == 0 and i != 0: print( f"test settings (main_orig eval): {args.type_feature} {args.test_mode} {args.feature_mode} {args.feature_num}" ) avg = logger.conditional_save_info(mode, average_meter, epoch) is_best = logger.rank_conditional_save_best(mode, avg, epoch) if is_best and not (mode == "train"): logger.save_img_comparison_as_best(mode, epoch) logger.conditional_summarize(mode, avg, is_best) if mode != "val": #if 1: helper.save_checkpoint({ # save checkpoint 'epoch': epoch, 'model': model.module.state_dict(), 'best_result': logger.best_result, 'optimizer': optimizer.state_dict(), 'args': args, }, is_best, epoch, logger.output_directory, args.type_feature, args.test_mode, args.feature_num, args.feature_mode, args.depth_adjust, i, every, "scratch") # draw features # run_info = [args.type_feature, alg_mode, feat_mode, model_orig] # if batch_data['rgb'] != None and 1 and (i % 1) == 0: # draw("sq", batch_data['rgb'], batch_data['d'], features, shape[1], run_info, i, result) return avg, is_best
def iterate(mode, args, loader, model, optimizer, logger, epoch): block_average_meter = AverageMeter() average_meter = AverageMeter() meters = [block_average_meter, average_meter] # switch to appropriate mode assert mode in ["train", "val", "eval", "test_prediction", "test_completion"], \ "unsupported mode: {}".format(mode) if mode == 'train': model.train() lr = helper.adjust_learning_rate(args.lr, optimizer, epoch) else: model.eval() lr = 0 torch.set_printoptions(profile="full") for i, batch_data in enumerate(loader): name = batch_data['name'][0] print(name) del batch_data['name'] print("i: ", i) # each batch data is 1 and has three keys d, gt, g and dim [1, 352, 1216] start = time.time() batch_data = { key: val.to(device) for key, val in batch_data.items() if val is not None } gt = batch_data[ 'gt'] if mode != 'test_prediction' and mode != 'test_completion' else None # if args.type_feature=="sq": # depth_adjustment(gt, False) data_time = time.time() - start start = time.time() if mode == "train": pred = model(batch_data) else: with torch.no_grad(): pred = model(batch_data) vis=False if vis: im = batch_data['gt'].detach().cpu().numpy() im_sq = im.squeeze() plt.figure() plt.imshow(im_sq) plt.show() # for i in range(im_sq.shape[0]): # print(f"{i} - {np.sum(im_sq[i])}") depth_loss, photometric_loss, smooth_loss, mask = 0, 0, 0, None if mode == 'train': # Loss 1: the direct depth supervision from ground truth label # mask=1 indicates that a pixel does not ground truth labels if 'sparse' in args.train_mode: depth_loss = depth_criterion(pred, batch_data['d']) print("d pts: ", len(torch.where(batch_data['d']>0)[0])) mask = (batch_data['d'] < 1e-3).float() elif 'dense' in args.train_mode: depth_loss = depth_criterion(pred, gt) mask = (gt < 1e-3).float() # Loss 2: the self-supervised photometric loss if args.use_pose: # create multi-scale pyramids pred_array = helper.multiscale(pred) rgb_curr_array = helper.multiscale(batch_data['rgb']) rgb_near_array = helper.multiscale(batch_data['rgb_near']) if mask is not None: mask_array = helper.multiscale(mask) num_scales = len(pred_array) # compute photometric loss at multiple scales for scale in range(len(pred_array)): pred_ = pred_array[scale] rgb_curr_ = rgb_curr_array[scale] rgb_near_ = rgb_near_array[scale] mask_ = None if mask is not None: mask_ = mask_array[scale] # compute the corresponding intrinsic parameters height_, width_ = pred_.size(2), pred_.size(3) intrinsics_ = kitti_intrinsics.scale(height_, width_) # inverse warp from a nearby frame to the current frame warped_ = homography_from(rgb_near_, pred_, batch_data['r_mat'], batch_data['t_vec'], intrinsics_) photometric_loss += photometric_criterion( rgb_curr_, warped_, mask_) * (2**(scale - num_scales)) # Loss 3: the depth smoothness loss smooth_loss = smoothness_criterion(pred) if args.w2 > 0 else 0 # backprop loss = depth_loss + args.w1 * photometric_loss + args.w2 * smooth_loss optimizer.zero_grad() loss.backward() zero_params(model) optimizer.step() gpu_time = time.time() - start # counting pixels in each bin #binned_pixels = np.load("value.npy", allow_pickle=True) #print(len(binned_pixels)) if (i % 1 == 0 and args.evaluate and args.instancewise) or\ (i % args.every == 0 and not args.evaluate and not args.instancewise): # global training # print(model.module.conv4[5].conv1.weight[0]) # print(model.conv4.5.bn2.weight) # print(model.module.parameter.grad) #print("*************swiches:") torch.set_printoptions(precision=7, sci_mode=False) if model.module.phi is not None: mmp = 1000 * model.module.phi phi = F.softplus(mmp) S = phi / torch.sum(phi) #print("S", S[1, -10:]) S_numpy= S.detach().cpu().numpy() if args.instancewise: global Ss if "Ss" not in globals(): Ss = [] Ss.append(S_numpy) else: Ss.append(S_numpy) # GLOBAL if (i % args.every ==0 and not args.evaluate and not args.instancewise and model.module.phi is not None): np.set_printoptions(precision=4) switches_2d_argsort = np.argsort(S_numpy, None) # 2d to 1d sort torch.Size([9, 31]) switches_2d_sort = np.sort(S_numpy, None) print("Switches: ") print(switches_2d_argsort[:10]) print(switches_2d_sort[:10]) print("and") print(switches_2d_argsort[-10:]) print(switches_2d_sort[-10:]) ##### saving global ranks global_ranks_path = lambda \ ii: f"ranks/{args.type_feature}/global/{folder_and_name[0]}/Ss_val_{folder_and_name[1]}_iter_{ii}.npy" global old_i if ("old_i" in globals()): print("old_i") if os.path.isfile(global_ranks_path(old_i)): os.remove(global_ranks_path(old_i)) folder_and_name = args.resume.split(os.sep)[-2:] os.makedirs(f"ranks/{args.type_feature}/global/{folder_and_name[0]}", exist_ok=True) np.save(global_ranks_path(i), S_numpy) old_i = i print("saving ranks") if args.type_feature == "sq": hor = switches_2d_argsort % S_numpy.shape[1] ver = np.floor(switches_2d_argsort // S_numpy.shape[1]) print(ver[:10],hor[:10]) print("and") print(ver[-10:], hor[-10:]) # measure accuracy and record loss with torch.no_grad(): mini_batch_size = next(iter(batch_data.values())).size(0) result = Result() if mode != 'test_prediction' and mode != 'test_completion': result.evaluate(pred.data, gt.data, photometric_loss) [ m.update(result, gpu_time, data_time, mini_batch_size) for m in meters ] logger.conditional_print(mode, i, epoch, lr, len(loader), block_average_meter, average_meter) logger.conditional_save_img_comparison(mode, i, batch_data, pred, epoch) logger.conditional_save_pred(mode, i, pred, epoch) draw=False if draw: ma = batch_data['rgb'].detach().cpu().numpy().squeeze() ma = np.transpose(ma, axes=[1, 2, 0]) # ma = np.uint8(ma) #ma2 = Image.fromarray(ma) ma2 = Image.fromarray(np.uint8(ma)).convert('RGB') # create rectangle image img1 = ImageDraw.Draw(ma2) if args.type_feature == "sq": size=40 print_square_num = 20 for ii in range(print_square_num): s_hor=hor[-ii].detach().cpu().numpy() s_ver=ver[-ii].detach().cpu().numpy() shape = [(s_hor * size, s_ver * size), ((s_hor + 1) * size, (s_ver + 1) * size)] img1.rectangle(shape, outline="red") tim = time.time() lala = ma2.save(f"switches_photos/squares/squares_{tim}.jpg") print("saving") elif args.type_feature == "lines": print_square_num = 20 r=1 parameter_mask = np.load("../kitti_pixels_to_lines.npy", allow_pickle=True) # for m in range(10,50): # im = Image.fromarray(parameter_mask[m]*155) # im = im.convert('1') # convert image to black and white # im.save(f"switches_photos/lala_{m}.jpg") for ii in range(print_square_num): points = parameter_mask[ii] y = np.where(points==1)[0] x = np.where(points == 1)[1] for p in range(len(x)): img1.ellipse((x[p] - r, y[p] - r, x[p] + r, y[p] + r), fill=(255, 0, 0, 0)) tim = time.time() lala = ma2.save(f"switches_photos/lines/lines_{tim}.jpg") print("saving") every = args.every if i % every ==0: print("saving") avg = logger.conditional_save_info(mode, average_meter, epoch) is_best = logger.rank_conditional_save_best(mode, avg, epoch) #is_best = True #saving all the checkpoints if is_best and not (mode == "train"): logger.save_img_comparison_as_best(mode, epoch) logger.conditional_summarize(mode, avg, is_best) if mode != "val": helper.save_checkpoint({ # save checkpoint 'epoch': epoch, 'model': model.module.state_dict(), 'best_result': logger.best_result, 'optimizer': optimizer.state_dict(), 'args': args, }, is_best, epoch, logger.output_directory, args.type_feature, i, every, qnet=True) if args.evaluate and args.instancewise: #filename = os.path.split(args.evaluate)[1] Ss_numpy = np.array(Ss) folder_and_name = args.evaluate.split(os.sep)[-3:] os.makedirs(f"ranks/instance/{folder_and_name[0]}", exist_ok=True) os.makedirs(f"ranks/instance/{folder_and_name[0]}/{folder_and_name[1]}", exist_ok=True) np.save(f"ranks/instance/{folder_and_name[0]}/{folder_and_name[1]}/Ss_val_{folder_and_name[2]}.npy", Ss) return avg, is_best
def iterate(mode, args, loader, model, optimizer, logger, epoch): block_average_meter = AverageMeter() average_meter = AverageMeter() meters = [block_average_meter, average_meter] # switch to appropriate mode assert mode in ["train", "val", "eval", "test_prediction", "test_completion"], \ "unsupported mode: {}".format(mode) encoder = model[0] decoder = model[1] if mode == 'train': encoder.train() decoder.train() lr = helper.adjust_learning_rate(args.lr, optimizer, epoch) else: encoder.train() decoder.train() lr = 0 torch.set_printoptions(profile="full") for i, batch_data in enumerate(loader): start = time.time() batch_data = { key: val.to(device) for key, val in batch_data.items() if val is not None } gt = batch_data[ 'gt'] if mode != 'test_prediction' and mode != 'test_completion' else None data_time = time.time() - start start = time.time() # torchvision.transforms.Resize # img = transform.resize(batch_data['rgb'], (192, 640)) # transform = transforms.Compose([ # transforms.Resize((round(192), round(640))), # # interpolation `BILINEAR` is applied by default # transforms.ToTensor(), # transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) # ]) # image = transform(batch_data['rgb']) # Otherwise, we only feed the image with frame_id 0 through the depth encoder features = encoder(batch_data['rgb']) outputs = decoder(features) # if i % 20 == 0: # print(outputs[('disp', 0)][1,0,:,40]) # pred = model(batch_data) # im = batch_data['d'].detach().cpu().numpy() # im_sq = im.squeeze() # plt.figure() # plt.imshow(im_sq) # plt.show() # for i in range(im_sq.shape[0]): # print(f"{i} - {np.sum(im_sq[i])}") pred = outputs[('disp', 0)] depth_loss, photometric_loss, smooth_loss, mask = 0, 0, 0, None if mode == 'train': # Loss 1: the direct depth supervision from ground truth label # mask=1 indicates that a pixel does not ground truth labels if 'sparse' in args.train_mode: depth_loss = depth_criterion(pred, batch_data['d']) mask = (batch_data['d'] < 1e-3).float() elif 'dense' in args.train_mode: depth_loss = depth_criterion(pred, gt) mask = (gt < 1e-3).float() # if i % 20 == 0: # print("\n\n\n gt \n") # print(gt[1,0,:,40]) # # Loss 2: the self-supervised photometric loss # if args.use_pose: # # create multi-scale pyramids # pred_array = helper.multiscale(pred) # rgb_curr_array = helper.multiscale(batch_data['rgb']) # rgb_near_array = helper.multiscale(batch_data['rgb_near']) # if mask is not None: # mask_array = helper.multiscale(mask) # num_scales = len(pred_array) # # # compute photometric loss at multiple scales # for scale in range(len(pred_array)): # pred_ = pred_array[scale] # rgb_curr_ = rgb_curr_array[scale] # rgb_near_ = rgb_near_array[scale] # mask_ = None # if mask is not None: # mask_ = mask_array[scale] # # # compute the corresponding intrinsic parameters # height_, width_ = pred_.size(2), pred_.size(3) # intrinsics_ = kitti_intrinsics.scale(height_, width_) # # # inverse warp from a nearby frame to the current frame # warped_ = homography_from(rgb_near_, pred_, # batch_data['r_mat'], # batch_data['t_vec'], intrinsics_) # photometric_loss += photometric_criterion( # rgb_curr_, warped_, mask_) * (2**(scale - num_scales)) # # Loss 3: the depth smoothness loss # smooth_loss = smoothness_criterion(pred) if args.w2 > 0 else 0 # backprop #loss = depth_loss + args.w1 * photometric_loss + args.w2 * smooth_loss loss = depth_loss optimizer.zero_grad() loss.backward() optimizer.step() gpu_time = time.time() - start if i % 50 == 0: print(i) # measure accuracy and record loss with torch.no_grad(): mini_batch_size = next(iter(batch_data.values())).size(0) result = Result() if mode != 'test_prediction' and mode != 'test_completion': result.evaluate(pred.data, gt.data, photometric_loss) [ m.update(result, gpu_time, data_time, mini_batch_size) for m in meters ] logger.conditional_print(mode, i, epoch, lr, len(loader), block_average_meter, average_meter) logger.conditional_save_img_comparison(mode, i, batch_data, pred, epoch) logger.conditional_save_pred(mode, i, pred, epoch) if i % 200 == 0: print(gpu_time) print("saving") avg = logger.conditional_save_info(mode, average_meter, epoch) is_best = logger.rank_conditional_save_best(mode, avg, epoch) if is_best and not (mode == "train"): logger.save_img_comparison_as_best(mode, epoch) logger.conditional_summarize(mode, avg, is_best) # helper.save_checkpoint({ # save checkpoint # 'epoch': epoch, # 'model': model.module.state_dict(), # 'best_result': logger.best_result, # 'optimizer': optimizer.state_dict(), # 'args': args, # }, is_best, epoch, logger.output_directory) return avg, is_best
def main(): global args, best_prec1 args = parser.parse_args() # create model if args.pretrained: print("=> using pre-trained model '{}'".format(args.arch)) else: print("=> creating model '{}'".format(args.arch)) if args.arch == 'alexnet': model = alexnet(pretrained=args.pretrained) elif args.arch == 'squeezenet1_0': model = squeezenet1_0(pretrained=args.pretrained) elif args.arch == 'squeezenet1_1': model = squeezenet1_1(pretrained=args.pretrained) elif args.arch == 'densenet121': model = densenet121(pretrained=args.pretrained) elif args.arch == 'densenet169': model = densenet169(pretrained=args.pretrained) elif args.arch == 'densenet201': model = densenet201(pretrained=args.pretrained) elif args.arch == 'densenet161': model = densenet161(pretrained=args.pretrained) elif args.arch == 'vgg11': model = vgg11(pretrained=args.pretrained) elif args.arch == 'vgg13': model = vgg13(pretrained=args.pretrained) elif args.arch == 'vgg16': model = vgg16(pretrained=args.pretrained) elif args.arch == 'vgg19': model = vgg19(pretrained=args.pretrained) elif args.arch == 'vgg11_bn': model = vgg11_bn(pretrained=args.pretrained) elif args.arch == 'vgg13_bn': model = vgg13_bn(pretrained=args.pretrained) elif args.arch == 'vgg16_bn': model = vgg16_bn(pretrained=args.pretrained) elif args.arch == 'vgg19_bn': model = vgg19_bn(pretrained=args.pretrained) elif args.arch == 'resnet18': model = resnet18(pretrained=args.pretrained) elif args.arch == 'resnet34': model = resnet34(pretrained=args.pretrained) elif args.arch == 'resnet50': model = resnet50(pretrained=args.pretrained) elif args.arch == 'resnet101': model = resnet101(pretrained=args.pretrained) elif args.arch == 'resnet152': model = resnet152(pretrained=args.pretrained) else: raise NotImplementedError # use cuda model.cuda() # model = torch.nn.parallel.DistributedDataParallel(model) # define loss and optimizer criterion = nn.CrossEntropyLoss().cuda() optimizer = optim.SGD(model.parameters(), lr=args.lr, momentum=args.momentum, weight_decay=args.weight_decay) # optionlly resume from a 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'] best_prec1 = checkpoint['best_prec1'] model.load_state_dict(checkpoint['state_dict']) optimizer.load_state_dict(checkpoint['optimizer']) print("=> loaded checkpoint '{}' (epoch {})".format( args.resume, checkpoint['epoch'])) else: print("=> no checkpoint found at '{}'".format(args.resume)) # cudnn.benchmark = True # Data loading train_loader, val_loader = data_loader(args.data, args.batch_size, args.workers, args.pin_memory) if args.evaluate: validate(val_loader, model, criterion, args.print_freq) return for epoch in range(args.start_epoch, args.epochs): adjust_learning_rate(optimizer, epoch, args.lr) # train for one epoch train(train_loader, model, criterion, optimizer, epoch, args.print_freq) # evaluate on validation set prec1, prec5 = validate(val_loader, model, criterion, args.print_freq) # remember the best prec@1 and save checkpoint is_best = prec1 > best_prec1 best_prec1 = max(prec1, best_prec1) save_checkpoint( { 'epoch': epoch + 1, 'arch': args.arch, 'state_dict': model.state_dict(), 'best_prec1': best_prec1, 'optimizer': optimizer.state_dict() }, is_best, args.arch + '.pth')
def train_model(args, model, dataset, writer=None, n_rounds=1, lth_pruner=None): root = args.exp_name + '/checkpoints/' criterion = nn.CrossEntropyLoss().cuda() optimizer = torch.optim.SGD(model.parameters(), args.lr, momentum=args.momentum, weight_decay=args.weight_decay) #Save initial weight files. init_weight_filename = args.exp_name + '/checkpoints/' + 'initial_state.pth.tar' helper.save_checkpoint(args, model, optimizer, init_weight_filename) for cur_round in range(n_rounds): best_acc = 0 for epoch in range(args.start_epoch, args.epochs): helper.adjust_learning_rate(optimizer, epoch, args) train_top1, train_top5, train_loss,model = trainer.train(dataset.train_loader,model,criterion,\ optimizer,epoch,args,lth_pruner,cur_round,mask_applied=args.mask_applied) val_top1, val_top5, val_loss = trainer.validate( dataset.test_loader, model, criterion, args) if writer is not None: writer.add_scalar("loss/train/" + str(cur_round), train_loss, epoch) writer.add_scalar("top1/train/" + str(cur_round), train_top1, epoch) writer.add_scalar("top5/train/" + str(cur_round), train_top5, epoch) writer.add_scalar("loss/val/" + str(cur_round), val_loss, epoch) writer.add_scalar("top1/val/" + str(cur_round), val_top1, epoch) writer.add_scalar("top5/val/" + str(cur_round), val_top5, epoch) if val_top1 >= best_acc: best_acc = val_top1 is_best = True filename = root + str(cur_round) + '_model_best.pth' helper.save_checkpoint(args, model, optimizer, filename) filename = root + str(cur_round) + '_current.pth' helper.save_checkpoint(args, model, optimizer, filename, epoch=epoch) filename = root + str(cur_round) + '_mask.pkl' if epoch in [0, 1, 2, 3]: #Save early epochs for late resetting. filename = root + 'epoch_' + str(epoch) + '_model.pth' helper.save_checkpoint(args, model, optimizer, filename, epoch=epoch)
def iterate(mode, args, loader, model, optimizer, logger, epoch): actual_epoch = epoch - args.start_epoch + args.start_epoch_bias block_average_meter = AverageMeter() block_average_meter.reset(False) average_meter = AverageMeter() meters = [block_average_meter, average_meter] # switch to appropriate mode assert mode in ["train", "val", "eval", "test_prediction", "test_completion"], \ "unsupported mode: {}".format(mode) if mode == 'train': model.train() lr = helper.adjust_learning_rate(args.lr, optimizer, actual_epoch, args) else: model.eval() lr = 0 torch.cuda.empty_cache() for i, batch_data in enumerate(loader): dstart = time.time() batch_data = { key: val.to(device) for key, val in batch_data.items() if val is not None } gt = batch_data[ 'gt'] if mode != 'test_prediction' and mode != 'test_completion' else None data_time = time.time() - dstart pred = None start = None gpu_time = 0 #start = time.time() #pred = model(batch_data) #gpu_time = time.time() - start #''' if (args.network_model == 'e'): start = time.time() st1_pred, st2_pred, pred = model(batch_data) else: start = time.time() pred = model(batch_data) if (args.evaluate): gpu_time = time.time() - start #''' depth_loss, photometric_loss, smooth_loss, mask = 0, 0, 0, None # inter loss_param st1_loss, st2_loss, loss = 0, 0, 0 w_st1, w_st2 = 0, 0 round1, round2, round3 = 1, 3, None if (actual_epoch <= round1): w_st1, w_st2 = 0.2, 0.2 elif (actual_epoch <= round2): w_st1, w_st2 = 0.05, 0.05 else: w_st1, w_st2 = 0, 0 if mode == 'train': # Loss 1: the direct depth supervision from ground truth label # mask=1 indicates that a pixel does not ground truth labels depth_loss = depth_criterion(pred, gt) if args.network_model == 'e': st1_loss = depth_criterion(st1_pred, gt) st2_loss = depth_criterion(st2_pred, gt) loss = (1 - w_st1 - w_st2 ) * depth_loss + w_st1 * st1_loss + w_st2 * st2_loss else: loss = depth_loss if i % multi_batch_size == 0: optimizer.zero_grad() loss.backward() if i % multi_batch_size == (multi_batch_size - 1) or i == (len(loader) - 1): optimizer.step() print("loss:", loss, " epoch:", epoch, " ", i, "/", len(loader)) if mode == "test_completion": str_i = str(i) path_i = str_i.zfill(10) + '.png' path = os.path.join(args.data_folder_save, path_i) vis_utils.save_depth_as_uint16png_upload(pred, path) if (not args.evaluate): gpu_time = time.time() - start # measure accuracy and record loss with torch.no_grad(): mini_batch_size = next(iter(batch_data.values())).size(0) result = Result() if mode != 'test_prediction' and mode != 'test_completion': result.evaluate(pred.data, gt.data, photometric_loss) [ m.update(result, gpu_time, data_time, mini_batch_size) for m in meters ] if mode != 'train': logger.conditional_print(mode, i, epoch, lr, len(loader), block_average_meter, average_meter) logger.conditional_save_img_comparison(mode, i, batch_data, pred, epoch) logger.conditional_save_pred(mode, i, pred, epoch) avg = logger.conditional_save_info(mode, average_meter, epoch) is_best = logger.rank_conditional_save_best(mode, avg, epoch) if is_best and not (mode == "train"): logger.save_img_comparison_as_best(mode, epoch) logger.conditional_summarize(mode, avg, is_best) return avg, is_best
def iterate(mode, args, loader, model, optimizer, logger, epoch): block_average_meter = AverageMeter() average_meter = AverageMeter() meters = [block_average_meter, average_meter] # switch to appropriate mode assert mode in ["train", "val", "eval", "test_prediction", "test_completion"], \ "unsupported mode: {}".format(mode) if mode == 'train': model.train() lr = helper.adjust_learning_rate(args.lr, optimizer, epoch) else: model.eval() lr = 0 for i, batch_data in enumerate(loader): #print(i) start = time.time() batch_data = { key: val.to(device) for key, val in batch_data.items() if val is not None } gt = batch_data[ 'gt'] if mode != 'test_prediction' and mode != 'test_completion' else None data_time = time.time() - start start = time.time() pred = model(batch_data) #im = batch_data['d'].detach().cpu().numpy() #im_sq = im.squeeze() #plt.figure() #plt.imshow(im_sq) #plt.show() #for i in range(im_sq.shape[0]): # print(f"{i} - {np.sum(im_sq[i])}") depth_loss, photometric_loss, smooth_loss, mask = 0, 0, 0, None if mode == 'train': # Loss 1: the direct depth supervision from ground truth label # mask=1 indicates that a pixel does not ground truth labels if 'sparse' in args.train_mode: depth_loss = depth_criterion(pred, batch_data['d']) mask = (batch_data['d'] < 1e-3).float() elif 'dense' in args.train_mode: depth_loss = depth_criterion(pred, gt) mask = (gt < 1e-3).float() # Loss 2: the self-supervised photometric loss if args.use_pose: # create multi-scale pyramids pred_array = helper.multiscale(pred) rgb_curr_array = helper.multiscale(batch_data['rgb']) rgb_near_array = helper.multiscale(batch_data['rgb_near']) if mask is not None: mask_array = helper.multiscale(mask) num_scales = len(pred_array) # compute photometric loss at multiple scales for scale in range(len(pred_array)): pred_ = pred_array[scale] rgb_curr_ = rgb_curr_array[scale] rgb_near_ = rgb_near_array[scale] mask_ = None if mask is not None: mask_ = mask_array[scale] # compute the corresponding intrinsic parameters height_, width_ = pred_.size(2), pred_.size(3) intrinsics_ = kitti_intrinsics.scale(height_, width_) # inverse warp from a nearby frame to the current frame warped_ = homography_from(rgb_near_, pred_, batch_data['r_mat'], batch_data['t_vec'], intrinsics_) photometric_loss += photometric_criterion( rgb_curr_, warped_, mask_) * (2**(scale - num_scales)) # Loss 3: the depth smoothness loss smooth_loss = smoothness_criterion(pred) if args.w2 > 0 else 0 # backprop loss = depth_loss + args.w1 * photometric_loss + args.w2 * smooth_loss optimizer.zero_grad() loss.backward() optimizer.step() gpu_time = time.time() - start if i % 50 ==0: # print(model.module.conv4[5].conv1.weight[0]) #print(model.conv4.5.bn2.weight) #print(model.module.parameter.grad) print(model.module.parameter) print(torch.argsort(model.module.parameter)) # measure accuracy and record loss with torch.no_grad(): mini_batch_size = next(iter(batch_data.values())).size(0) result = Result() if mode != 'test_prediction' and mode != 'test_completion': result.evaluate(pred.data, gt.data, photometric_loss) [ m.update(result, gpu_time, data_time, mini_batch_size) for m in meters ] logger.conditional_print(mode, i, epoch, lr, len(loader), block_average_meter, average_meter) logger.conditional_save_img_comparison(mode, i, batch_data, pred, epoch) logger.conditional_save_pred(mode, i, pred, epoch) if i % 20 == 0: avg = logger.conditional_save_info(mode, average_meter, epoch) is_best = logger.rank_conditional_save_best(mode, avg, epoch) if is_best and not (mode == "train"): logger.save_img_comparison_as_best(mode, epoch) logger.conditional_summarize(mode, avg, is_best) helper.save_checkpoint({ # save checkpoint 'epoch': epoch, 'model': model.module.state_dict(), 'best_result': logger.best_result, 'optimizer': optimizer.state_dict(), 'args': args, }, is_best, epoch, logger.output_directory) return avg, is_best
def iterate(mode, args, loader, model, optimizer, logger, epoch): block_average_meter = AverageMeter() average_meter = AverageMeter() meters = [block_average_meter, average_meter] block_average_meter_intensity = AverageIntensity() average_meter_intensity = AverageIntensity() meters_intensity = [block_average_meter_intensity, average_meter_intensity] block_average_meter_intensity_pure = AverageIntensity() average_meter_intensity_pure = AverageIntensity() meters_intensity_pure = [ block_average_meter_intensity_pure, average_meter_intensity_pure ] # switch to appropriate mode assert mode in ["train", "val", "eval", "test_prediction", "test_completion"], \ "unsupported mode: {}".format(mode) if mode == 'train': model.train() # lr = args.lr #helper.adjust_learning_rate(args.lr, optimizer, epoch) if args.lradj > 0: lr = helper.adjust_learning_rate(args.lr, optimizer, epoch, args.lradj) else: lr = args.lr else: model.eval() lr = 0 global batch_num for i, batch_data in enumerate(loader): if mode == "train": batch_num += 1 start = time.time() file_name = batch_data['filename'] batch_data = { key: val.cuda() for key, val in batch_data.items() if (key != "filename" and val is not None) } # batch_data = {key:val.cuda() for key,val in batch_data.items() if val is not None} gt = batch_data[ 'gt'] if mode != 'test_prediction' and mode != 'test_completion' else None # Ireal as gt_intensity gt_intensity = batch_data[ 'gt_intensity'] if mode != 'test_prediction' and mode != 'test_completion' else None gt_intensity_pure = batch_data[ 'gt_intensity_pure'] if mode != 'test_prediction' and mode != 'test_completion' else None data_time = time.time() - start start = time.time() pred = model(batch_data) pred_intensity_pure = pred[0][:, 1, :, :].unsqueeze(1) pred_intensity = pred[1] pred = pred[0][:, 0, :, :].unsqueeze(1) depth_loss, intensity_loss, photometric_loss, smooth_loss, pure_loss, mask = 0, 0, 0, 0, 0, None if mode == 'train': # Loss 1: the direct depth supervision from ground truth label # mask=1 indicates that a pixel does not ground truth labels if 'sparse' in args.train_mode: depth_loss = depth_criterion(pred, batch_data['d']) mask = (batch_data['d'] < 1e-3).float() elif 'dense' in args.train_mode: depth_loss = depth_criterion(pred, gt) pure_loss = intensity_pure_criterion(pred_intensity_pure, gt_intensity_pure) intensity_loss = intensity_criterion(pred_intensity, gt_intensity) mask = (gt < 1e-3).float() # Loss 2: the self-supervised photometric loss if args.use_pose: # create multi-scale pyramids pred_array = helper.multiscale(pred) rgb_curr_array = helper.multiscale(batch_data['rgb']) rgb_near_array = helper.multiscale(batch_data['rgb_near']) if mask is not None: mask_array = helper.multiscale(mask) num_scales = len(pred_array) # compute photometric loss at multiple scales for scale in range(len(pred_array)): pred_ = pred_array[scale] rgb_curr_ = rgb_curr_array[scale] rgb_near_ = rgb_near_array[scale] mask_ = None if mask is not None: mask_ = mask_array[scale] # compute the corresponding intrinsic parameters height_, width_ = pred_.size(2), pred_.size(3) intrinsics_ = kitti_intrinsics.scale(height_, width_) # inverse warp from a nearby frame to the current frame warped_ = homography_from(rgb_near_, pred_, batch_data['r_mat'], batch_data['t_vec'], intrinsics_) photometric_loss += photometric_criterion( rgb_curr_, warped_, mask_) * (2**(scale - num_scales)) # Loss 3: the depth smoothness loss smooth_loss = smoothness_criterion(pred) if args.w2 > 0 else 0 # backprop loss = depth_loss + args.wi * intensity_loss + args.wpure * pure_loss # loss = depth_loss + wi * intensity_loss + args.w1*photometric_loss + args.w2*smooth_loss optimizer.zero_grad() loss.backward() optimizer.step() #log train if batch_num % 25 == 0 and mode == "train": logger.writer.add_scalar('train/loss_total', loss, batch_num) logger.writer.add_scalar('train/loss_depth', depth_loss, batch_num) logger.writer.add_scalar('train/loss_intensity', intensity_loss, batch_num) logger.writer.add_scalar('train/loss_pure', pure_loss, batch_num) logger.writer.add_scalar('train/lr', lr, batch_num) #if batch_num % 200 == 0: # logger.writer.add_image('train/Ipure', pred_intensity_pure[0], batch_num) gpu_time = time.time() - start # measure accuracy and record loss with torch.no_grad(): mini_batch_size = next(iter(batch_data.values())).size(0) result = Result() # result_intensity = Result() result_intensity_pure = Result_intensity() result_intensity = Result_intensity() if mode != 'test_prediction' and mode != 'test_completion': # result.evaluate(pred.data, gt.data, photometric_loss) result.evaluate(pred.data, gt.data, photometric_loss) result_intensity_pure.evaluate(pred_intensity_pure.data, gt_intensity_pure.data) result_intensity.evaluate(pred_intensity.data, gt_intensity.data) [ m.update(result, gpu_time, data_time, mini_batch_size) for m in meters ] [ m_intensity.update(result_intensity, gpu_time, data_time, mini_batch_size) for m_intensity in meters_intensity ] [ m_intensity_pure.update(result_intensity_pure, gpu_time, data_time, mini_batch_size) for m_intensity_pure in meters_intensity_pure ] logger.conditional_print(mode, i, epoch, lr, len(loader), block_average_meter, average_meter) # logger.conditional_save_img_comparison(mode, i, batch_data, pred, epoch) logger.conditional_save_pred_named(mode, file_name[0], pred, epoch) logger.conditional_print(mode, i, epoch, lr, len(loader), block_average_meter_intensity_pure, average_meter_intensity_pure, "Ipure") logger.conditional_print(mode, i, epoch, lr, len(loader), block_average_meter_intensity, average_meter_intensity, "Intensity") # logger.conditional_save_img_comparison_with_intensity(mode, i, batch_data, pred, pred_intensity, epoch) logger.conditional_save_img_comparison_with_intensity2( mode, i, batch_data, pred, pred_intensity_pure, pred_intensity, epoch) # run when eval, bt always = 1 #logger.conditional_save_pred_named_with_intensity(mode, file_name[0], pred, pred_intensity, epoch) avg = logger.conditional_save_info(mode, average_meter, epoch) avg_intensity = logger.conditional_save_info_intensity( mode, average_meter_intensity, epoch) # is_best = logger.rank_conditional_save_best(mode, avg, epoch) is_best = logger.rank_conditional_save_best_with_intensity( mode, avg, avg_intensity, epoch) if is_best and not (mode == "train"): logger.save_img_comparison_as_best(mode, epoch) logger.conditional_summarize(mode, avg, is_best) logger.conditional_summarize_intensity(mode, avg_intensity) return avg, avg_intensity, is_best
def main(): global args, best_prec1 global cur_itrs args = parser.parse_args() print(args.mode) # STEP1: model if args.mode=='baseline_train': model = initialize_model(use_resnet=True, pretrained=False, nclasses=200) elif args.mode=='pretrain': model = deeplab_network.deeplabv3_resnet50(num_classes=args.num_classes, output_stride=args.output_stride, pretrained_backbone=False) set_bn_momentum(model.backbone, momentum=0.01) elif args.mode=='finetune': model = initialize_model(use_resnet=True, pretrained=False, nclasses=3) # load the pretrained model if args.pretrained_model: if os.path.isfile(args.pretrained_model): print("=> loading pretrained model '{}'".format(args.pretrained_model)) checkpoint = torch.load(args.pretrained_model) args.start_epoch = checkpoint['epoch'] best_prec1 = checkpoint['best_prec1'] model.load_state_dict(checkpoint['state_dict']) optimizer.load_state_dict(checkpoint['optimizer']) print("=> loaded pretrained model '{}' (epoch {})".format(args.resume, checkpoint['epoch'])) if torch.cuda.is_available: model = model.cuda() # STEP2: criterion and optimizer if args.mode in ['baseline_train', 'finetune']: criterion = nn.CrossEntropyLoss() optimizer = torch.optim.SGD(model.parameters(), lr=args.lr, momentum=args.momentum, weight_decay=args.weight_decay) # train_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=args.step_size, gamma=args.gamma) elif args.mode=='pretrain': criterion = nn.MSELoss() optimizer = torch.optim.SGD(params=[ {'params': model.backbone.parameters(), 'lr': 0.1*args.lr}, {'params': model.classifier.parameters(), 'lr': args.lr}, ], lr=args.lr, momentum=0.9, weight_decay=args.weight_decay) scheduler = PolyLR(optimizer, args.total_itrs, power=0.9) # STEP3: loss/prec record if args.mode in ['baseline_train', 'finetune']: train_losses = [] train_top1s = [] train_top5s = [] test_losses = [] test_top1s = [] test_top5s = [] elif args.mode == 'pretrain': train_losses = [] test_losses = [] # STEP4: optionlly resume from a checkpoint if args.resume: print('resume') if os.path.isfile(args.resume): print("=> loading checkpoint '{}'".format(args.resume)) if args.mode in ['baseline_train', 'finetune']: checkpoint = torch.load(args.resume) args.start_epoch = checkpoint['epoch'] best_prec1 = checkpoint['best_prec1'] model.load_state_dict(checkpoint['state_dict']) optimizer.load_state_dict(checkpoint['optimizer']) datafile = args.resume.split('.pth')[0] + '.npz' load_data = np.load(datafile) train_losses = list(load_data['train_losses']) train_top1s = list(load_data['train_top1s']) train_top5s = list(load_data['train_top5s']) test_losses = list(load_data['test_losses']) test_top1s = list(load_data['test_top1s']) test_top5s = list(load_data['test_top5s']) elif args.mode=='pretrain': checkpoint = torch.load(args.resume) args.start_epoch = checkpoint['epoch'] model.load_state_dict(checkpoint['state_dict']) optimizer.load_state_dict(checkpoint['optimizer']) scheduler.load_state_dict(checkpoint['scheduler']) cur_itrs = checkpoint['cur_itrs'] datafile = args.resume.split('.pth')[0] + '.npz' load_data = np.load(datafile) train_losses = list(load_data['train_losses']) # test_losses = list(load_data['test_losses']) print("=> loaded checkpoint '{}' (epoch {})".format(args.resume, checkpoint['epoch'])) else: print("=> no checkpoint found at '{}'".format(args.resume)) # STEP5: train! if args.mode in ['baseline_train', 'finetune']: # data from utils import TinyImageNet_data_loader print('color_distortion:', color_distortion) train_loader, val_loader = TinyImageNet_data_loader(args.dataset, args.batch_size,color_distortion=args.color_distortion) # if evaluate the model if args.evaluate: print('evaluate this model on validation dataset') validate(val_loader, model, criterion, args.print_freq) return for epoch in range(args.start_epoch, args.epochs): adjust_learning_rate(optimizer, epoch, args.lr) time1 = time.time() #timekeeping # train for one epoch model.train() loss, top1, top5 = train(train_loader, model, criterion, optimizer, epoch, args.print_freq) train_losses.append(loss) train_top1s.append(top1) train_top5s.append(top5) # evaluate on validation set model.eval() loss, prec1, prec5 = validate(val_loader, model, criterion, args.print_freq) test_losses.append(loss) test_top1s.append(prec1) test_top5s.append(prec5) # remember the best prec@1 and save checkpoint is_best = prec1 > best_prec1 best_prec1 = max(prec1, best_prec1) save_checkpoint({ 'epoch': epoch + 1, 'mode': args.mode, 'state_dict': model.state_dict(), 'best_prec1': best_prec1, 'optimizer': optimizer.state_dict() }, is_best, args.mode + '_' + args.dataset +'.pth') np.savez(args.mode + '_' + args.dataset +'.npz', train_losses=train_losses,train_top1s=train_top1s,train_top5s=train_top5s, test_losses=test_losses,test_top1s=test_top1s, test_top5s=test_top5s) # np.savez(args.mode + '_' + args.dataset +'.npz', train_losses=train_losses) time2 = time.time() #timekeeping print('Elapsed time for epoch:',time2 - time1,'s') print('ETA of completion:',(time2 - time1)*(args.epochs - epoch - 1)/60,'minutes') print() elif args.mode=='pretrain': #data from utils import TinyImageNet_data_loader # args.dataset = 'tiny-imagenet-200' args.batch_size = 16 train_loader, val_loader = TinyImageNet_data_loader(args.dataset, args.batch_size, col=True) # if evaluate the model, show some results if args.evaluate: print('evaluate this model on validation dataset') visulization(val_loader, model, args.start_epoch) return # for epoch in range(args.start_epoch, args.epochs): epoch = 0 while True: if cur_itrs >= args.total_itrs: return # adjust_learning_rate(optimizer, epoch, args.lr) time1 = time.time() #timekeeping model.train() # train for one epoch # loss, _, _ = train(train_loader, model, criterion, optimizer, epoch, args.print_freq, colorization=True,scheduler=scheduler) # train_losses.append(loss) # model.eval() # # evaluate on validation set # loss, _, _ = validate(val_loader, model, criterion, args.print_freq, colorization=True) # test_losses.append(loss) save_checkpoint({ 'epoch': epoch + 1, 'mode': args.mode, 'state_dict': model.state_dict(), 'optimizer': optimizer.state_dict(), 'scheduler':scheduler.state_dict(), "cur_itrs": cur_itrs }, True, args.mode + '_' + args.dataset +'.pth') np.savez(args.mode + '_' + args.dataset +'.npz', train_losses=train_losses) # scheduler.step() time2 = time.time() #timekeeping print('Elapsed time for epoch:',time2 - time1,'s') print('ETA of completion:',(time2 - time1)*(args.total_itrs - cur_itrs - 1)/60,'minutes') print() epoch += 1
def main(args): if args.debug: import pdb; pdb.set_trace(); tb_dir = args.exp_name+'/tb_logs/' ckpt_dir = args.exp_name + '/checkpoints/' if not os.path.exists(args.exp_name): os.mkdir(args.exp_name) os.mkdir(tb_dir) os.mkdir(ckpt_dir) #writer = SummaryWriter(tb_dir+'{}'.format(args.exp_name), flush_secs=10) writer = SummaryWriter(tb_dir, flush_secs=10) # create model print("=> creating model: ") os.system('nvidia-smi') #model = models.__dict__[args.arch]() #model = resnet_dilated.Resnet18_32s(num_classes=21) print(args.no_pre_train,' pretrain') #model = resnet18_fcn.Resnet18_fcn(num_classes=args.n_classes,pre_train=args.no_pre_train) model_map = { 'deeplabv3_resnet18': arma_network.deeplabv3_resnet18, 'deeplabv3_resnet50': arma_network.deeplabv3_resnet50, 'fcn_resnet18': arma_network.fcn_resnet18, #'deeplabv3_resnet101': network.deeplabv3_resnet101, # 'deeplabv3plus_resnet18': network.deeplabv3plus_resnet18, # 'deeplabv3plus_resnet50': network.deeplabv3plus_resnet50, # 'deeplabv3plus_resnet101': network.deeplabv3plus_resnet101 } model = model_map['deeplabv3_resnet50'](arma=False,num_classes=args.n_classes) model = model.cuda() model = nn.DataParallel(model) optimizer = torch.optim.SGD(model.parameters(), args.lr, momentum=args.momentum, weight_decay=args.weight_decay) # optionally resume from a checkpoint if args.resume: if os.path.isfile(args.resume): model,optimizer,args = helper.load_checkpoint(args,model,optimizer) else: print("=> no checkpoint found at '{}'".format(args.resume)) #USE this only when batch size is fixed. #This takes time, but optimizes to crazy speeds once input is fixed. cudnn.benchmark = True #Load dataloaders augmentations = aug.Compose([aug.RandomCrop(512),aug.RandomHorizontallyFlip(5),\ aug.RandomRotate(30),aug.RandomSizedCrop(512)]) my_dataset = pascalVOCLoader(args=args,root=args.data,sbd_path=args.data,\ augmentations=augmentations) my_dataset.get_loaders() init_weight_filename ='initial_state.pth.tar' helper.save_checkpoint(args,model,optimizer,custom_name=init_weight_filename) with open(args.exp_name+'/'+'args.pkl','wb') as fout: pickle.dump(args,fout) best_iou = -100.0 for epoch in range(args.start_epoch, args.epochs): helper.adjust_learning_rate(optimizer, epoch, args) train_loss = trainer.train(my_dataset.train_loader,model,optimizer,epoch,args,writer) val_loss,scores,class_iou,running_metrics_val = trainer.validate(my_dataset.val_loader, model,epoch,args,writer) if scores["Mean IoU : \t"] >= best_iou: best_iou = scores["Mean IoU : \t"] is_best = True if not args.multiprocessing_distributed or (args.multiprocessing_distributed and args.rank % ngpus_per_node == 0): if epoch in [0,1,2,3,4,5,6,7,8]: helper.save_checkpoint(args,model,optimizer,epoch,custom_name=str(epoch)+'.pth') if args.save_freq is None: helper.save_checkpoint(args,model,optimizer,epoch,is_best=is_best,periodic=False) else: helper.save_checkpoint(args,model,optimizer,epoch,is_best=is_best,periodic=True) with open(args.exp_name+'/running_metric.pkl','wb') as fout: pickle.dump(running_metrics_val,fout)
def main_worker(gpu, ngpus_per_node, args): args.gpu = gpu if args.debug: import pdb; pdb.set_trace(); if not args.multiprocessing_distributed or (args.multiprocessing_distributed and args.rank % ngpus_per_node == 0): tb_dir = args.exp_name+'/tb_logs/' ckpt_dir = args.exp_name + '/checkpoints/' if not os.path.exists(args.exp_name): os.mkdir(args.exp_name) os.mkdir(tb_dir) os.mkdir(ckpt_dir) print("writing to : ",tb_dir+'{}'.format(args.exp_name),args.rank,ngpus_per_node) #writer = SummaryWriter(tb_dir+'{}'.format(args.exp_name), flush_secs=10) writer = SummaryWriter(tb_dir, flush_secs=10) # suppress printing if not master if args.multiprocessing_distributed and args.gpu != 0: def print_pass(*args): pass builtins.print = print_pass if args.gpu is not None: print("Use GPU: {} for training".format(args.gpu)) if args.distributed: if args.dist_url == "env://" and args.rank == -1: args.rank = int(os.environ["RANK"]) if args.multiprocessing_distributed: args.rank = args.rank * ngpus_per_node + gpu dist.init_process_group(backend=args.dist_backend, init_method=args.dist_url, world_size=args.world_size, rank=args.rank) # create model print("=> creating model: ") #model = models.__dict__[args.arch]() model = resnet_dilated.Resnet18_32s(num_classes=21) if args.distributed: print("distributed") if args.gpu is not None: torch.cuda.set_device(args.gpu) model.cuda(args.gpu) args.batch_size = int(args.batch_size / ngpus_per_node) args.workers = int((args.workers + ngpus_per_node - 1) / ngpus_per_node) model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.gpu]) else: model.cuda() model = torch.nn.parallel.DistributedDataParallel(model) elif args.gpu is not None: torch.cuda.set_device(args.gpu) model = model.cuda(args.gpu) # comment out the following line for debugging raise NotImplementedError("Only DistributedDataParallel is supported.") else: raise NotImplementedError("Only DistributedDataParallel is supported.") optimizer = torch.optim.SGD(model.parameters(), args.lr, momentum=args.momentum, weight_decay=args.weight_decay) # optionally resume from a checkpoint if args.resume: if os.path.isfile(args.resume): model,optimizer,args = helper.load_checkpoint(args,model,optimizer) else: print("=> no checkpoint found at '{}'".format(args.resume)) #USE this only when batch size is fixed. #This takes time, but optimizes to crazy speeds once input is fixed. cudnn.benchmark = True #Load dataloaders augmentations = aug.Compose([aug.RandomCrop(512),aug.RandomHorizontallyFlip(5),aug.RandomRotate(30),aug.RandomSizedCrop(512)]) my_dataset = pascalVOCLoader(args=args,root='/scratch0/shishira/pascal_voc/',sbd_path='/scratch0/shishira/pascal_voc/',\ augmentations=augmentations) my_dataset.get_loaders() init_weight_filename ='initial_state.pth.tar' helper.save_checkpoint(args,model,optimizer,custom_name=init_weight_filename) with open(args.exp_name+'/'+'args.pkl','wb') as fout: pickle.dump(args,fout) best_iou = -100.0 for epoch in range(args.start_epoch, args.epochs): if args.distributed: my_dataset.train_sampler.set_epoch(epoch) helper.adjust_learning_rate(optimizer, epoch, args) train_loss = trainer.train(my_dataset.train_loader,model,optimizer,epoch,args,writer) val_loss,scores,class_iou = trainer.validate(my_dataset.val_loader, model,epoch,args,writer) if scores["Mean IoU : \t"] >= best_iou: best_iou = scores["Mean IoU : \t"] is_best = True if not args.multiprocessing_distributed or (args.multiprocessing_distributed and args.rank % ngpus_per_node == 0): if epoch in [0,1,2,3,4,5,6,7,8]: helper.save_checkpoint(args,model,optimizer,epoch,custom_name=str(epoch)+'.pth') if args.save_freq is None: helper.save_checkpoint(args,model,optimizer,epoch,is_best=is_best,periodic=False) else: helper.save_checkpoint(args,model,optimizer,epoch,is_best=is_best,periodic=True)
def main(): global args, best_prec1 args = parser.parse_args() os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu #Log log_format = '[%(asctime)s] %(message)s' logging.basicConfig(stream=sys.stdout, level=logging.INFO, format=log_format, datefmt='%d %I:%M:%S') t = time.time() local_time = time.localtime(t) if not os.path.exists('./log'): os.mkdir('./log') fh = logging.FileHandler( os.path.join('log/train-{}{}{:02}{}'.format(args.arch, local_time.tm_year % 2000, local_time.tm_mon, t))) fh.setFormatter(logging.Formatter(log_format)) logging.getLogger().addHandler(fh) # create model if args.pretrained: print("=> using pre-trained model '{}'".format(args.arch)) else: print("=> creating model '{}'".format(args.arch)) if args.arch == 'alexnet': model = alexnet(pretrained=args.pretrained) elif args.arch == 'alexnet_s': model = alexnet_s(pretrained=args.pretrained) elif args.arch == 'alexnet_s_addchannel': model = alexnet_s_addchannel(pretrained=args.pretrained) elif args.arch == 'alexnet_s_addchannel_fullshuffle': model = alexnet_s_addchannel_fullshuffle(pretrained=args.pretrained) elif args.arch == 'alexnet_s_addchannelx2_fullshuffle': model = alexnet_s_addchannelx2_fullshuffle(pretrained=args.pretrained) elif args.arch == 'alexnet_s1_2x': model = alexnet_s1_2x(pretrained=args.pretrained) elif args.arch == 'squeezenet1_0': model = squeezenet1_0(pretrained=args.pretrained) elif args.arch == 'squeezenet1_1': model = squeezenet1_1(pretrained=args.pretrained) elif args.arch == 'densenet121': model = densenet121(pretrained=args.pretrained) elif args.arch == 'densenet169': model = densenet169(pretrained=args.pretrained) elif args.arch == 'densenet201': model = densenet201(pretrained=args.pretrained) elif args.arch == 'densenet161': model = densenet161(pretrained=args.pretrained) elif args.arch == 'vgg11': model = vgg11(pretrained=args.pretrained) elif args.arch == 'vgg13': model = vgg13(pretrained=args.pretrained) elif args.arch == 'vgg16': model = vgg16(pretrained=args.pretrained) elif args.arch == 'vgg19': model = vgg19(pretrained=args.pretrained) elif args.arch == 'vgg11_bn': model = vgg11_bn(pretrained=args.pretrained) elif args.arch == 'vgg13_bn': model = vgg13_bn(pretrained=args.pretrained) elif args.arch == 'vgg16_bn': model = vgg16_bn(pretrained=args.pretrained) elif args.arch == 'vgg19_bn': model = vgg19_bn(pretrained=args.pretrained) elif args.arch == 'vgg16_bn_s3_addchannel': model = vgg16_bn_s3_addchannel(pretrained=args.pretrained) elif args.arch == 'vgg16_bn_s3_addchannel_v2': model = vgg16_bn_s3_addchannel_v2(pretrained=args.pretrained) elif args.arch == 'resnet18': model = resnet18(pretrained=args.pretrained) elif args.arch == 'resnet34': model = resnet34(pretrained=args.pretrained) elif args.arch == 'resnet34_s2_addchannel': model = resnet34_s2_addchannel(pretrained=args.pretrained) elif args.arch == 'resnet34_s3_x15': model = resnet34_s3_x15(pretrained=args.pretrained) elif args.arch == 'resnet34_s3_x175': model = resnet34_s3_x175(pretrained=args.pretrained) elif args.arch == 'resnet50': model = resnet50(pretrained=args.pretrained) elif args.arch == 'resnet50_s': model = resnet50_s(pretrained=args.pretrained) elif args.arch == 'resnet50_block_s': model = resnet50_block_s(pretrained=args.pretrained) elif args.arch == 'resnet50_block_s_addchannel': model = resnet50_block_s_addchannel(pretrained=args.pretrained) elif args.arch == 'resnet50_s3_addchannel': model = resnet50_s3_addchannel(pretrained=args.pretrained) elif args.arch == 'resnet50_s3_addchannel175': model = resnet50_s3_addchannel175(pretrained=args.pretrained) elif args.arch == 'resnet50_justaddgroup': model = resnet50_justaddgroup(pretrained=args.pretrained) elif args.arch == 'resnet101': model = resnet101(pretrained=args.pretrained) elif args.arch == 'resnet152': model = resnet152(pretrained=args.pretrained) else: raise NotImplementedError # use cuda model.cuda() # model = torch.nn.parallel.DistributedDataParallel(model) # define loss and optimizer criterion = nn.CrossEntropyLoss().cuda() optimizer = optim.SGD(model.parameters(), lr=args.lr, momentum=args.momentum, weight_decay=args.weight_decay) # optionlly resume from a 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'] best_prec1 = checkpoint['best_prec1'] model.load_state_dict(checkpoint['state_dict']) optimizer.load_state_dict(checkpoint['optimizer']) print("=> loaded checkpoint '{}' (epoch {})".format( args.resume, checkpoint['epoch'])) else: print("=> no checkpoint found at '{}'".format(args.resume)) # cudnn.benchmark = True # Data loading train_loader, val_loader = data_loader(args.data, args.batch_size, args.workers, args.pin_memory) if args.evaluate: validate(val_loader, model, criterion, args.print_freq) return for epoch in range(args.start_epoch, args.epochs): adjust_learning_rate(optimizer, epoch, args.lr) # train for one epoch train(train_loader, model, criterion, optimizer, epoch, args.print_freq) # evaluate on validation set prec1, prec5 = validate(val_loader, model, criterion, args.print_freq) # remember the best prec@1 and save checkpoint is_best = prec1 > best_prec1 best_prec1 = max(prec1, best_prec1) save_checkpoint( { 'epoch': epoch + 1, 'arch': args.arch, 'state_dict': model.state_dict(), 'best_prec1': best_prec1, 'optimizer': optimizer.state_dict() }, is_best, args.arch + '.pth')
def iterate(mode, args, loader, model, optimizer, logger, epoch): block_average_meter = AverageMeter() average_meter = AverageMeter() meters = [block_average_meter, average_meter] # switch to appropriate mode assert mode in ["train", "val", "eval", "test_prediction", "test_completion"], \ "unsupported mode: {}".format(mode) if mode == 'train': model.train() lr = helper.adjust_learning_rate(args.lr, optimizer, epoch) else: model.eval() lr = 0 print("\nTraining") prune_type = "sq" # sq, vlines, nothing square_choice = args.training_sparse_opt if prune_type == "sq": print(f"Features: squares\n Square choice: {square_choice}") for i, batch_data in enumerate(loader): start = time.time() batch_data = { key: val.to(device) for key, val in batch_data.items() if val is not None } gt = batch_data[ 'gt'] if mode != 'test_prediction' and mode != 'test_completion' else None data_time = time.time() - start start = time.time() if prune_type == "vlines": np.random.seed(10) lines_unmasked = np.random.choice(352, 20, replace=False) lines_unmasked = np.arange(352) lines_all = np.arange(352) lines_masked = [x for x in lines_all if x not in lines_unmasked] batch_data['d'][:, :, lines_masked] = 0 print(batch_data['d'].shape) print("lines unmasked", lines_unmasked) elif prune_type == "sq": A = np.load("ranks/switches_2D_equal_iter_390.npy", allow_pickle=True) # with np.printoptions(precision=5): # print("switches", A) # get the ver and hor coordinates of the most important squares A_2d_argsort = np.argsort(A, None)[::-1] if square_choice == "most": squares_top_file = "ranks/sq/global/squares_most.npy" A_2d_argsort = np.load(squares_top_file)[::-1] ver = np.floor(A_2d_argsort // A.shape[1]) hor = A_2d_argsort % A.shape[1] A_list = np.stack([ver, hor]).transpose() square_size = 40 squares_top_num = 50 if square_choice == "full": squares_top = A_list if square_choice == "most": squares_top = A_list[:squares_top_num] if square_choice == "best_sw": squares_top = A_list[:squares_top_num] if square_choice == "latin_sw": # creating latin grid (with big squares/blocks) hor_large = np.linspace(0, 30, 7) ver_larger = np.arange(10) all_squares = np.arange(len(A_list)) bins_2d_latin = binned_statistic_2d( ver, hor, all_squares, 'min', bins=[ver_larger, hor_large]) bins_2d_latin.statistic best_latin = bins_2d_latin.statistic[-3:].flatten().astype(int) best_latin_coors = list(A_list[best_latin]) for i1 in A_list: elem_in = False # check if the block already contains a small square for i2 in best_latin_coors: if i1[0] == i2[0] and i1[1] == i2[1]: elem_in = True if not elem_in: best_latin_coors.append(i1) if len(best_latin_coors) == squares_top_num: break squares_top = np.array(best_latin_coors) elif square_choice == "latin": np.random.seed(12) squares_latin_evenlyspaced = [] # create blocks, choose k random blocks and have fixed first block hor_large = np.linspace(0, 30, 7) ver_large = np.arange(10) # random sample from positive blocks hor_large_rand = np.random.choice(len(hor_large), squares_top_num) ver_large_rand = np.random.choice([6, 7, 8], squares_top_num) # selecting a small square from A_list with given corrdinates within a block for j in range(len(hor_large_rand)): elem = \ np.where((A_list[:, 0] == ver_large_rand[j]) & (A_list[:, 1] == hor_large[hor_large_rand[j]]))[0][0] squares_latin_evenlyspaced.append(elem) squares_top = A_list[squares_latin_evenlyspaced] elif square_choice == "random_all": np.random.seed(12) rand_idx = np.random.choice(len(A_list), squares_top_num) print(rand_idx) squares_top = A_list[rand_idx] elif square_choice == "random_pos": # from squares which include depth points np.random.seed(12) # choose from the squares which have roughly positive number of depth points rand_idx = np.random.choice(len(A_list[:93]), squares_top_num) print(rand_idx) squares_top = A_list[rand_idx] # after selecting indices of the squares save in squares_top squares_top_scaled = np.array(squares_top) * square_size mask = np.zeros((352, 1216)) bin_ver = np.arange(0, 352, square_size) bin_ver = np.append(bin_ver, oheight) bin_hor = np.arange(0, 1216, square_size) bin_hor = np.append(bin_hor, owidth) # filling in the mask with selected squares up to squares_top_num (e.g. 20) # print("Number of squares selected: ", len(squares_top)) # print(squares_top) for it in range( len(squares_top) ): # in all but full should be equal to squares_top_num ver = int(squares_top[it][0]) hor = int(squares_top[it][1]) # print("ver", bin_ver[ver], bin_ver[ver+1], "hor", bin_hor[hor], bin_hor[hor+1] ) mask[bin_ver[ver]:bin_ver[ver + 1], bin_hor[hor]:bin_hor[hor + 1]] = 1 aaa1 = batch_data['d'].detach().cpu().numpy() batch_data['d'] = torch.einsum( "abcd, cd->abcd", [batch_data['d'], torch.Tensor(mask).to(device)]) aaa2 = batch_data['d'].detach().cpu().numpy() # # # from PIL import Image # # img = Image.fromarray(aaa[0, :, :, :], 'RGB') # # #img.save('my.png') # # img.show() pred = model(batch_data) # im = batch_data['d'].detach().cpu().numpy() # im_sq = im.squeeze() # plt.figure() # plt.imshow(im_sq) # plt.show() # for i in range(im_sq.shape[0]): # print(f"{i} - {np.sum(im_sq[i])}") depth_loss, photometric_loss, smooth_loss, mask = 0, 0, 0, None if mode == 'train': # Loss 1: the direct depth supervision from ground truth label # mask=1 indicates that a pixel does not ground truth labels if 'sparse' in args.train_mode: depth_loss = depth_criterion(pred, batch_data['d']) mask = (batch_data['d'] < 1e-3).float() elif 'dense' in args.train_mode: depth_loss = depth_criterion(pred, gt) mask = (gt < 1e-3).float() # Loss 2: the self-supervised photometric loss if args.use_pose: # create multi-scale pyramids pred_array = helper.multiscale(pred) rgb_curr_array = helper.multiscale(batch_data['rgb']) rgb_near_array = helper.multiscale(batch_data['rgb_near']) if mask is not None: mask_array = helper.multiscale(mask) num_scales = len(pred_array) # compute photometric loss at multiple scales for scale in range(len(pred_array)): pred_ = pred_array[scale] rgb_curr_ = rgb_curr_array[scale] rgb_near_ = rgb_near_array[scale] mask_ = None if mask is not None: mask_ = mask_array[scale] # compute the corresponding intrinsic parameters height_, width_ = pred_.size(2), pred_.size(3) intrinsics_ = kitti_intrinsics.scale(height_, width_) # inverse warp from a nearby frame to the current frame warped_ = homography_from(rgb_near_, pred_, batch_data['r_mat'], batch_data['t_vec'], intrinsics_) photometric_loss += photometric_criterion( rgb_curr_, warped_, mask_) * (2**(scale - num_scales)) # Loss 3: the depth smoothness loss smooth_loss = smoothness_criterion(pred) if args.w2 > 0 else 0 # backprop loss = depth_loss + args.w1 * photometric_loss + args.w2 * smooth_loss optimizer.zero_grad() loss.backward() optimizer.step() gpu_time = time.time() - start # measure accuracy and record loss with torch.no_grad(): mini_batch_size = next(iter(batch_data.values())).size(0) result = Result() if mode != 'test_prediction' and mode != 'test_completion': result.evaluate(pred.data, gt.data, photometric_loss) [ m.update(result, gpu_time, data_time, mini_batch_size) for m in meters ] logger.conditional_print(mode, i, epoch, lr, len(loader), block_average_meter, average_meter) logger.conditional_save_img_comparison(mode, i, batch_data, pred, epoch) logger.conditional_save_pred(mode, i, pred, epoch) every = 100 if i % 500 == 0: # every 100 batches/images (before it was after the entire dataset - two tabs/on if statement) avg = logger.conditional_save_info(mode, average_meter, epoch) is_best = logger.rank_conditional_save_best(mode, avg, epoch) if is_best and not (mode == "train"): logger.save_img_comparison_as_best(mode, epoch) logger.conditional_summarize(mode, avg, is_best) helper.save_checkpoint({ # save checkpoint 'epoch': epoch, 'model': model.module.state_dict(), 'best_result': logger.best_result, 'optimizer': optimizer.state_dict(), 'args': args, }, is_best, epoch, logger.output_directory, args.type_feature, i, every, "scratch") return avg, is_best