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 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 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")