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)
