def resume(self, cfg): """ load model and optimizer state """ if cfg.schedule.resume is not None: load_path = cfg.schedule.resume else: load_path = os.path.join(cfg.save_dir, 'model_last.pth') checkpoint = torch.load(load_path, map_location=lambda storage, loc: storage) self.logger.log('loaded {}, epoch {}'.format(load_path, checkpoint['epoch'])) if hasattr(self.model, 'module'): load_model_weight(self.model.module, checkpoint, self.logger) else: load_model_weight(self.model, checkpoint, self.logger) if 'optimizer' in checkpoint: self.optimizer.load_state_dict(checkpoint['optimizer']) self.epoch = checkpoint['epoch'] + 1 self.logger.log('resumed at epoch: {}'.format(self.epoch)) if 'iter' in checkpoint: self._iter = checkpoint['iter'] + 1 self.logger.log('resumed at steps: {}'.format(self._iter)) else: self.logger.log('No optimizer parameters in checkpoint.')
def main(config, model_path: str, output_path: str, input_shape=(320, 320)): logger = Logger(local_rank=-1, save_dir=config.save_dir, use_tensorboard=False) # Create model and load weights model = build_model(config.model) checkpoint = torch.load(model_path, map_location=lambda storage, loc: storage) load_model_weight(model, checkpoint, logger) # Convert backbone weights for RepVGG models if config.model.arch.backbone.name == "RepVGG": deploy_config = config.model deploy_config.arch.backbone.update({"deploy": True}) deploy_model = build_model(deploy_config) from nanodet.model.backbone.repvgg import repvgg_det_model_convert model = repvgg_det_model_convert(model, deploy_model) # TorchScript: tracing the model with dummy inputs with torch.no_grad(): dummy_input = torch.zeros(1, 3, input_shape[0], input_shape[1]) # Batch size = 1 model.eval().cpu() model_traced = torch.jit.trace(model, example_inputs=dummy_input).eval() model_traced.save(output_path) print("Finished export to TorchScript")
def __init__(self, cfg, model_path, logger, device='cuda:0'): self.cfg = cfg self.device = device model = build_model(cfg.model) ckpt = torch.load(model_path, map_location=lambda storage, loc: storage) load_model_weight(model, ckpt, logger) self.model = model.to(device).eval() self.pipeline = Pipeline(cfg.data.val.pipeline, cfg.data.val.keep_ratio)
def load_model(self, cfg): load_path = cfg.schedule.load_model checkpoint = torch.load(load_path, map_location=lambda storage, loc: storage) self.logger.log('loaded {}, epoch {}'.format(load_path, checkpoint['epoch'])) if hasattr(self.model, 'module'): load_model_weight(self.model.module, checkpoint, self.logger) else: load_model_weight(self.model, checkpoint, self.logger)
def main(config, model_path, output_path, input_shape=(320, 320)): logger = Logger(-1, config.save_dir, False) model = build_model(config.model) checkpoint = torch.load(model_path, map_location=lambda storage, loc: storage) load_model_weight(model, checkpoint, logger) dummy_input = torch.autograd.Variable(torch.randn(1, 3, input_shape[0], input_shape[1])) torch.onnx.export(model, dummy_input, output_path, verbose=True, keep_initializers_as_inputs=True, opset_version=11) print('finished exporting onnx ')
def main(config, model_path, output_path, input_shape=(320, 320)): logger = Logger(-1, config.save_dir, False) model = build_model(config.model) checkpoint = torch.load(model_path, map_location=lambda storage, loc: storage) load_model_weight(model, checkpoint, logger) if config.model.arch.backbone.name == 'RepVGG': deploy_config = config.model deploy_config.arch.backbone.update({'deploy': True}) deploy_model = build_model(deploy_config) from nanodet.model.backbone.repvgg import repvgg_det_model_convert model = repvgg_det_model_convert(model, deploy_model) dummy_input = torch.autograd.Variable( torch.randn(1, 3, input_shape[0], input_shape[1])) dynamic_axes = { "input": { 0: "batch_size" }, "output1": { 0: "batch_size" }, "output2": { 0: "batch_size" }, "output3": { 0: "batch_size" }, "output4": { 0: "batch_size" }, "output5": { 0: "batch_size" }, "output6": { 0: "batch_size" } } input_names = ['input'] output_names = [ 'output1', 'output2', 'output3', 'output4', 'output5', 'output6' ] torch.onnx.export(model, dummy_input, output_path, verbose=True, keep_initializers_as_inputs=True, opset_version=12, input_names=input_names, output_names=output_names) import onnx from onnxsim import simplify model = onnx.load(output_path) # convert model model_simp, check = simplify(model) onnx.save(model_simp, output_path) print('finished exporting onnx ')
def main(config, model_path, output_path, input_shape=(320, 320)): logger = Logger(-1, config.save_dir, False) model = build_model(config.model) checkpoint = torch.load(model_path, map_location=lambda storage, loc: storage) load_model_weight(model, checkpoint, logger) if config.model.arch.backbone.name == 'RepVGG': deploy_config = config.model deploy_config.arch.backbone.update({'deploy': True}) deploy_model = build_model(deploy_config) from nanodet.model.backbone.repvgg import repvgg_det_model_convert model = repvgg_det_model_convert(model, deploy_model) dummy_input = torch.autograd.Variable(torch.randn(1, 3, input_shape[0], input_shape[1])) torch.onnx.export(model, dummy_input, output_path, verbose=True, keep_initializers_as_inputs=True, opset_version=11) print('finished exporting onnx ')
def __init__(self, cfg, model_path, logger, device='cuda:0'): self.cfg = cfg self.device = device model = build_model(cfg.model) ckpt = torch.load(model_path, map_location=lambda storage, loc: storage) load_model_weight(model, ckpt, logger) if cfg.model.arch.backbone.name == 'RepVGG': deploy_config = cfg.model deploy_config.arch.backbone.update({'deploy': True}) deploy_model = build_model(deploy_config) from nanodet.model.backbone.repvgg import repvgg_det_model_convert model = repvgg_det_model_convert(model, deploy_model) self.model = model.to(device).eval() self.pipeline = Pipeline(cfg.data.val.pipeline, cfg.data.val.keep_ratio)
def __init__(self, model_path, cfg_path, *args, **kwargs): from nanodet.model.arch import build_model from nanodet.util import Logger, cfg, load_config, load_model_weight super(NanoDetTorch, self).__init__(*args, **kwargs) print(f'Using PyTorch as inference backend') print(f'Using weight: {model_path}') # load model self.model_path = model_path self.cfg_path = cfg_path load_config(cfg, cfg_path) self.logger = Logger(-1, cfg.save_dir, False) self.model = build_model(cfg.model) checkpoint = self.torch.load(model_path, map_location=lambda storage, loc: storage) load_model_weight(self.model, checkpoint, self.logger)
def export_ONNX_model(self, config, model_path, output_path, logger, input_shape=(320, 320)): model = build_model(config.model) checkpoint = torch.load(model_path, map_location=lambda storage, loc: storage) load_model_weight(model, checkpoint, logger) dummy_input = torch.autograd.Variable( torch.randn(1, 3, input_shape[0], input_shape[1])) torch.onnx.export(model, dummy_input, output_path, verbose=True, keep_initializers_as_inputs=True, opset_version=11)
def __init__(self, cfg, model_path, logger, device='cuda:0'): """ :param cfg: :param model_path: :param logger: :param device: """ self.cfg = cfg self.device = device model = build_model(cfg.model) ckpt = torch.load(model_path, map_location=lambda storage, loc: storage) load_model_weight(model, ckpt, logger) print('INFO: {:s} loaded.'.format(model_path)) self.model = model.to(device).eval() self.pipeline = Pipeline(cfg.data.val.pipeline, cfg.data.val.keep_ratio)
def resume(self, cfg): """ load model and optimizer state """ if cfg.schedule.resume is not None: load_path = cfg.schedule.resume else: load_path = os.path.join(cfg.save_dir, 'model_last.pth') model_path = load_path['load_model'] if not os.path.isfile(model_path): self.logger.log('[Err]: invalid checkpoint path.') return checkpoint = torch.load(model_path, map_location=lambda storage, loc: storage) self.logger.log('loaded {}, epoch {}'.format(load_path, checkpoint['epoch'])) print('loaded {}, epoch {}'.format(load_path, checkpoint['epoch'])) if hasattr(self.model, 'module'): load_model_weight(self.model.module, checkpoint, self.logger) else: load_model_weight(self.model, checkpoint, self.logger) if 'optimizer' in checkpoint: self.optimizer.load_state_dict(checkpoint['optimizer']) # print(self.optimizer.param_groups[0]['lr']) self.optimizer.param_groups[0][ 'lr'] = self.cfg.schedule.optimizer.lr print('Current lr: ', self.optimizer.param_groups[0]['lr']) self.epoch = checkpoint['epoch'] + 1 self.logger.log('resumed at epoch: {}'.format(self.epoch)) if 'iter' in checkpoint: self._iter = checkpoint['iter'] + 1 self.logger.log('resumed at steps: {}'.format(self._iter)) else: self.logger.log('No optimizer parameters in checkpoint.')
def main(config, model_path, output_path, input_shape=(320, 320), batch_size=1): logger = Logger(-1, config.save_dir, False) model = build_model(config.model) checkpoint = torch.load(model_path, map_location=lambda storage, loc: storage) load_model_weight(model, checkpoint, logger) dummy_input = torch.autograd.Variable( torch.randn(batch_size, 3, input_shape[0], input_shape[1])) torch.onnx.export(model, dummy_input, output_path, verbose=True, keep_initializers_as_inputs=True, opset_version=12) onnx_model = onnx.load(output_path) # load onnx model model_simp, check = simplify(onnx_model) assert check, "Simplified ONNX model could not be validated" onnx.save(model_simp, output_path) print('finished exporting onnx ')
def main(config, model_path, output_path, input_shape=(320, 320)): logger = Logger(-1, config.save_dir, False) model = build_model(config.model) checkpoint = torch.load(model_path, map_location=lambda storage, loc: storage) load_model_weight(model, checkpoint, logger) if config.model.arch.backbone.name == "RepVGG": deploy_config = config.model deploy_config.arch.backbone.update({"deploy": True}) deploy_model = build_model(deploy_config) from nanodet.model.backbone.repvgg import repvgg_det_model_convert model = repvgg_det_model_convert(model, deploy_model) dummy_input = torch.autograd.Variable( torch.randn(1, 3, input_shape[0], input_shape[1]) ) torch.onnx.export( model, dummy_input, output_path, verbose=True, keep_initializers_as_inputs=True, opset_version=11, input_names=["data"], output_names=["output"], ) logger.log("finished exporting onnx ") logger.log("start simplifying onnx ") input_data = {"data": dummy_input.detach().cpu().numpy()} model_sim, flag = onnxsim.simplify(output_path, input_data=input_data) if flag: onnx.save(model_sim, output_path) logger.log("simplify onnx successfully") else: logger.log("simplify onnx failed")
def main(args): load_config(cfg, args.config) if cfg.model.arch.head.num_classes != len(cfg.class_names): raise ValueError( "cfg.model.arch.head.num_classes must equal len(cfg.class_names), " "but got {} and {}".format(cfg.model.arch.head.num_classes, len(cfg.class_names))) local_rank = int(args.local_rank) torch.backends.cudnn.enabled = True torch.backends.cudnn.benchmark = True mkdir(local_rank, cfg.save_dir) logger = NanoDetLightningLogger(cfg.save_dir) logger.dump_cfg(cfg) if args.seed is not None: logger.info("Set random seed to {}".format(args.seed)) pl.seed_everything(args.seed) logger.info("Setting up data...") train_dataset = build_dataset(cfg.data.train, "train") val_dataset = build_dataset(cfg.data.val, "test") evaluator = build_evaluator(cfg.evaluator, val_dataset) train_dataloader = torch.utils.data.DataLoader( train_dataset, batch_size=cfg.device.batchsize_per_gpu, shuffle=True, num_workers=cfg.device.workers_per_gpu, pin_memory=True, collate_fn=naive_collate, drop_last=True, ) val_dataloader = torch.utils.data.DataLoader( val_dataset, batch_size=cfg.device.batchsize_per_gpu, shuffle=False, num_workers=cfg.device.workers_per_gpu, pin_memory=True, collate_fn=naive_collate, drop_last=False, ) logger.info("Creating model...") task = TrainingTask(cfg, evaluator) if "load_model" in cfg.schedule: ckpt = torch.load(cfg.schedule.load_model) if "pytorch-lightning_version" not in ckpt: warnings.warn( "Warning! Old .pth checkpoint is deprecated. " "Convert the checkpoint with tools/convert_old_checkpoint.py ") ckpt = convert_old_model(ckpt) load_model_weight(task.model, ckpt, logger) logger.info("Loaded model weight from {}".format( cfg.schedule.load_model)) model_resume_path = (os.path.join(cfg.save_dir, "model_last.ckpt") if "resume" in cfg.schedule else None) accelerator = None if len(cfg.device.gpu_ids) <= 1 else "ddp" trainer = pl.Trainer( default_root_dir=cfg.save_dir, max_epochs=cfg.schedule.total_epochs, gpus=cfg.device.gpu_ids, check_val_every_n_epoch=cfg.schedule.val_intervals, accelerator=accelerator, log_every_n_steps=cfg.log.interval, num_sanity_val_steps=0, resume_from_checkpoint=model_resume_path, callbacks=[ProgressBar(refresh_rate=0)], # disable tqdm bar logger=logger, benchmark=True, gradient_clip_val=cfg.get("grad_clip", 0.0), ) trainer.fit(task, train_dataloader, val_dataloader)