def main(cfg): orig_cwd = hydra.utils.get_original_cwd() print(f" Evaluating the checkpoint " f"'{cfg.test.trained_ckpt_for_inference}'") loaded_struct = torch.load( pjn(orig_cwd, cfg.test.trained_ckpt_for_inference)) cfg.train.train_val_split = 1.0 # ^ associate all of the train data with the train_loader below # (do not split the train data into train + validation) train_loader, _, test_loader = train.get_dataloaders(cfg, (1, 0, 1)) interval_bounds, label2count_list = make_label2count_list(cfg) model = SDCNet(label2count_list, cfg.model.supervised, load_pretr_weights_vgg=False) model.load_state_dict(loaded_struct['model_state_dict'], strict=True) additional_cfg = {'device': None} if not cfg.resources.disable_cuda and torch.cuda.is_available(): additional_cfg['device'] = torch.device('cuda') model = model.cuda() else: additional_cfg['device'] = torch.device('cpu') optimizer = None trainer = train.TrainManager( model, optimizer, cfg, additional_cfg, train_loader=None, val_loader=None, ) print() datadir = pjn(cfg.dataset.dataset_rootdir, f"part_{cfg.dataset.part}") print(f" Evaluating on the (whole) train data and on the test data " f"(in '{datadir}')") mae_train, mse_train = trainer.validate(train_loader) print(f" Metrics on the (whole) train data: " f"MAE: {mae_train:.2f}, MSE: {mse_train:.2f}") mae_test, mse_test = trainer.validate(test_loader) print(f" Metrics on the test data: " f"MAE: {mae_test:.2f}, MSE: {mse_test:.2f}") if cfg.test.visualize: vis_dir_name = f"visualized_part_{cfg.dataset.part}_test_set_predictions" vis_dir_print = pjn(os.path.relpath(os.getcwd(), orig_cwd), vis_dir_name) print( f" Visualized predictions are being saved to '{vis_dir_print}':") visualize_predictions(cfg, model, test_loader, vis_dir_name)
def main(cfg): orig_cwd = hydra.utils.get_original_cwd() print(f" Exporting the checkpoint " f"'{cfg.test.trained_ckpt_for_inference}'") ckpt_path = pjn(orig_cwd, cfg.test.trained_ckpt_for_inference) loaded_struct = torch.load(ckpt_path) interval_bounds, label2count_list = make_label2count_list(cfg) model = SDCNet(label2count_list, cfg.model.supervised, load_pretr_weights_vgg=False) model.load_state_dict(loaded_struct['model_state_dict'], strict=True) batch_size = 1 x = torch.randn(batch_size, 3, 64 * 1, 64 * 1, requires_grad=False) if not cfg.resources.disable_cuda and torch.cuda.is_available(): model = model.cuda() x = x.cuda() p1, ext = os.path.splitext(ckpt_path) dir_, bname = os.path.split(p1) try: torch.onnx.export(model, x, p1 + ".onnx", opset_version=13) except: print(" (!) Failed to export the checkpoint to ONNX format") else: print(" Successfully exported the checkpoint to ONNX format") try: traced_script_module = torch.jit.trace(model, x) traced_script_module.save(bname + "_jit_trace.pt") except: print( " (!) Failed to export the checkpoint to 'torch jit trace' format" ) else: print( " Successfully exported the checkpoint to 'torch jit trace' format" ) try: script_module = torch.jit.script(model) script_module.save(bname + "_jit_script.pt") except: print( " (!) Failed to export the checkpoint to 'torch jit script' format" ) else: print( " Successfully exported the checkpoint to 'torch jit script' format" )
def main(cfg): """ Create data loaders, the model instance, optimizer and TrainManager() object. Run the training process. """ orig_cwd = hydra.utils.get_original_cwd() train_loader, val_loader, test_loader = get_dataloaders(cfg, (1, 1, 0)) interval_bounds, label2count_list = make_label2count_list(cfg) model = SDCNet( label2count_list, cfg.model.supervised, load_pretr_weights_vgg=True) if cfg.train.pretrained_ckpt: print(f" Using pretrained model and its checkpoint " f"'{cfg.train.pretrained_ckpt}'") loaded_struct = torch.load(pjn(orig_cwd, cfg.train.pretrained_ckpt)) model.load_state_dict(loaded_struct['model_state_dict'], strict=True) additional_cfg = {'device': None} if not cfg.resources.disable_cuda and torch.cuda.is_available(): additional_cfg['device'] = torch.device('cuda') model = model.cuda() else: additional_cfg['device'] = torch.device('cpu') optimizer = torch.optim.SGD( model.parameters(), lr=cfg.train.lr_schedule.lr_init, momentum=cfg.train.optimizer.momentum, weight_decay=cfg.train.optimizer.weight_decay) additional_cfg['tbx_wrtr_dir'] = os.getcwd() with SummaryWriter(additional_cfg['tbx_wrtr_dir']) as tbx_wrtr: additional_cfg['tbx_wrtr'] = tbx_wrtr trainer = TrainManager( model, optimizer, cfg, additional_cfg, train_loader=train_loader, val_loader=val_loader) trainer.train()
def main(cfg): orig_cwd = hydra.utils.get_original_cwd() print(f" Running inference using checkpoint " f"'{cfg.test.trained_ckpt_for_inference}'") ckpt_path = pjn(orig_cwd, cfg.test.trained_ckpt_for_inference) loaded_struct = torch.load(ckpt_path) interval_bounds, label2count_list = make_label2count_list(cfg) model = SDCNet(label2count_list, cfg.model.supervised, load_pretr_weights_vgg=False) model.load_state_dict(loaded_struct['model_state_dict'], strict=True) if not cfg.resources.disable_cuda and torch.cuda.is_available(): model = model.cuda() if cfg.test.visualize: vis_dir_print = pjn(os.path.relpath(os.getcwd(), orig_cwd), 'visualized_predictions') print(f" Visualized predictions are being saved to '{vis_dir_print}'") get_predictions(cfg, model)
class Object(object): pass class Cfg: # imitates config from hydra (only 3 parameters from the full config # are required here to instantiate SDCNet) def __init__(self, num_intervals, interval_step, partition_method): self.dataset = Object() self.dataset.num_intervals = num_intervals self.model = Object() self.model.interval_step = interval_step self.model.partition_method = partition_method cfg = Cfg(22, 0.5, 2) # for part_A #cfg = Cfg(7, 0.5, 2) # for part_B interval_bounds, label2count_list = make_label2count_list(cfg) # ================================================================ device = torch.device("cuda" if torch.cuda.is_available() else "cpu") sdcnet_instance = SDCNet(label2count_list, supervised=True, load_pretr_weights_vgg=True).to(device) batch_size = 1 h = 64 * 3 w = 64 * 2 summary(sdcnet_instance, (3, h, w)) x = torch.randn(batch_size, 3, h, w, requires_grad=False).to(device) out_list = sdcnet_instance(x)
def get_dataloaders(cfg, train_val_test_mask): """ Constructs Dataset objects and corresponding DataLoader objects. Args: cfg: the global configuration (hydra). train_val_test_mask: a list or tuple containing three values that will be interpreted as booleans, for example (1,1,0). The corresponding DataLoaders for the (train, val, test) filtered list of Datasets will be created and returned. Returns: The tuple (train_loader, val_loader, test_loader). The loaders corresponding to the `False` values in `train_val_test_mask` are set to `None`. """ interval_bounds, label2count_list = make_label2count_list(cfg) assert len(label2count_list) == interval_bounds.shape[0] + 1, \ "Number of class labels must be equal to (1 + number of interval " \ "boundaries), but the equality does not hold" rgb_mean_train = dtst.calc_rgb_mean_train(cfg) composed_transf_train = transforms.Compose([ dtst.Normalize(rgb_mean_train), dtst.ToCHW(), dtst.RandomHorizontalFlip(p=0.5), dtst.QuasiRandomCrop(), dtst.PadToMultipleOf64(), dtst.AddGtCountsLabels(interval_bounds), ]) composed_transf_test = transforms.Compose([ dtst.Normalize(rgb_mean_train), dtst.ToCHW(), dtst.PadToMultipleOf64(), ]) if train_val_test_mask[0]: train_dataset = dtst.ShanghaiTechDataset( cfg, subdir='train_data', shuffle_seed=42, rel_inds=(0.0, cfg.train.train_val_split), transform=composed_transf_train) train_loader = DataLoader( train_dataset, batch_size=cfg.train.batch_size, shuffle=True, num_workers=4) else: train_loader = None if train_val_test_mask[1]: val_dataset = dtst.ShanghaiTechDataset( cfg, subdir='train_data', shuffle_seed=42, rel_inds=(cfg.train.train_val_split, 1.0), transform=composed_transf_test) val_loader = DataLoader( val_dataset, batch_size=cfg.validation.batch_size, shuffle=False, num_workers=4) else: val_loader = None if train_val_test_mask[2]: test_dataset = dtst.ShanghaiTechDataset( cfg, subdir='test_data', shuffle_seed=None, rel_inds=(0.0, 1.0), transform=composed_transf_test) test_loader = DataLoader( test_dataset, batch_size=cfg.test.batch_size, shuffle=False, num_workers=4) else: test_loader = None return train_loader, val_loader, test_loader
def get_dataloaders(cfg, train_val_test_mask): interval_bounds, label2count_list = make_label2count_list(cfg) assert len(label2count_list) == interval_bounds.shape[0] + 1, \ "Number of class labels must be equal to (1 + number of interval " \ "boundaries), but the equality does not hold" rgb_mean_train = dtst.calc_rgb_mean_train(cfg) composed_transf_train = transforms.Compose([ dtst.Normalize(rgb_mean_train), dtst.ToCHW(), dtst.RandomHorizontalFlip(p=0.5), dtst.QuasiRandomCrop(), dtst.PadToMultipleOf64(), dtst.AddGtCountsLabels(interval_bounds), ]) composed_transf_test = transforms.Compose([ dtst.Normalize(rgb_mean_train), dtst.ToCHW(), dtst.PadToMultipleOf64(), ]) if train_val_test_mask[0]: train_dataset = dtst.ShanghaiTechDataset( cfg, subdir='train_data', shuffle_seed=42, rel_inds=(0.0, cfg.train.train_val_split), transform=composed_transf_train) train_loader = DataLoader( train_dataset, batch_size=cfg.train.batch_size, shuffle=True, num_workers=4) else: train_loader = None if train_val_test_mask[1]: val_dataset = dtst.ShanghaiTechDataset( cfg, subdir='train_data', shuffle_seed=42, rel_inds=(cfg.train.train_val_split, 1.0), transform=composed_transf_test) val_loader = DataLoader( train_dataset, batch_size=cfg.validation.batch_size, shuffle=False, num_workers=4) else: val_loader = None if train_val_test_mask[2]: test_dataset = dtst.ShanghaiTechDataset( cfg, subdir='test_data', shuffle_seed=None, rel_inds=(0.0, 1.0), transform=composed_transf_test) test_loader = DataLoader( test_dataset, batch_size=cfg.test.batch_size, shuffle=False, num_workers=4) else: test_loader = None return train_loader, val_loader, test_loader