Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add resnet-ssd #53

Open
wants to merge 16 commits into
base: main
Choose a base branch
from
Prev Previous commit
Next Next commit
[Feature] quant deploy
[Fix] quant deploy bug
  • Loading branch information
yongyang committed Jun 28, 2022
commit 89da4c7d109523d93e094084723ac84091f0da4f
1 change: 1 addition & 0 deletions up/commands/__init__.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
from .train import Train # noqa
from .inference import Inference # noqa
from .eval import Eval # noqa
from .quant_deploy import QuantDeploy # noqa
from .flops import Flops # noqa
61 changes: 61 additions & 0 deletions up/commands/quant_deploy.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,61 @@
from __future__ import division

# Standard Library
import argparse


# Import from local
from .subcommand import Subcommand
from up.utils.general.yaml_loader import load_yaml # IncludeLoader
from up.utils.general.registry_factory import SUBCOMMAND_REGISTRY, DEPLOY_REGISTRY
from up.utils.general.user_analysis_helper import send_info
from up.utils.general.global_flag import QUANT_FLAG

__all__ = ['QuantDeploy']


@SUBCOMMAND_REGISTRY.register('quant_deploy')
class QuantDeploy(Subcommand):
def add_subparser(self, name, parser):
sub_parser = parser.add_parser(name,
description='subcommand for deploying',
help='deploy a model')
sub_parser.add_argument('--config',
dest='config',
required=True,
help='settings of detection in yaml format')
sub_parser.add_argument('--cfg_type',
dest='cfg_type',
type=str,
default='up',
help='config type (up or pod)')
sub_parser.add_argument('--opts',
help='options to replace yaml config',
default=None,
nargs=argparse.REMAINDER)
sub_parser.set_defaults(run=_main)
return sub_parser


def main(args):
cfg = load_yaml(args.config, args.cfg_type)
cfg['args'] = {
'opts': args.opts
}
cfg['runtime'] = cfg.setdefault('runtime', {})
runner_cfg = cfg['runtime'].get('runner', {})
runner_cfg['type'] = runner_cfg.get('type', 'base')
runner_cfg['kwargs'] = runner_cfg.get('kwargs', {})
cfg['runtime']['runner'] = runner_cfg
QUANT_FLAG.flag = True

send_info(cfg, func="quant_deploy")
if runner_cfg['type'] == "quant":
quant_deploy = DEPLOY_REGISTRY.get("quant")(cfg)
quant_deploy.deploy()
else:
print("Need quant in cfg yaml.")


def _main(args):
main(args)
2 changes: 1 addition & 1 deletion up/tasks/cls/models/heads/cls_head.py
Original file line number Diff line number Diff line change
@@ -63,7 +63,7 @@ def forward_net(self, x):
x = self.get_pool_output(x)
x = self.get_dropout(x)
logits = self.get_logits(x)
return {'logits': logits}
return {'logits': logits, 'deploy_output_node': logits}

def forward(self, input):
return self.forward_net(input)
1 change: 1 addition & 0 deletions up/tasks/quant/__init__.py
Original file line number Diff line number Diff line change
@@ -1,2 +1,3 @@
from .deploy import * # noqa
from .runner import * # noqa
from .models import * # noqa
4 changes: 4 additions & 0 deletions up/tasks/quant/deploy/__init__.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,4 @@
try: # noqa
from .quant_deploy import QuantDeploy # noqa
except: # noqa
pass # noqa
56 changes: 56 additions & 0 deletions up/tasks/quant/deploy/quant_deploy.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,56 @@
import torch
from up.utils.general.log_helper import default_logger as logger
from up.utils.general.registry_factory import DEPLOY_REGISTRY
from up.tasks.quant.runner import QuantRunner
from mqbench.utils.state import enable_quantization
from up.utils.general.global_flag import DEPLOY_FLAG


__all__ = ['QuantDeploy']


@DEPLOY_REGISTRY.register('quant')
class QuantDeploy(QuantRunner):
def __init__(self, config, work_dir='./'):
self.config = config
super(QuantDeploy, self).__init__(config, work_dir, False)

def build(self):
self.build_ema()
self.build_saver()
self.load_ckpt()
self.build_model()
self.get_onnx_dummy_input()
self.quantize_model()
self.resume_model_from_quant()
enable_quantization(self.model)

def resume_model_from_quant(self):
self.model.load_state_dict(self.ckpt['model'])

def get_onnx_dummy_input(self):
self.model.cuda().eval()
DEPLOY_FLAG.flag = False
self.build_dataloaders()
self.build_hooks()
batch = self.get_batch('test')
output = self.model(batch)
self.dummy_input = {k: v for k, v in output.items() if torch.is_tensor(v)}
DEPLOY_FLAG.flag = True

def deploy(self):
logger.info("deploy model")
from mqbench.convert_deploy import convert_deploy
deploy_backend = self.config['quant']['deploy_backend']
self.model.cuda().eval()
print('ONNX input shape is: ', self.dummy_input['image'].shape)

for index, mname in enumerate(self.model_list):
mod = getattr(self.model, mname)
print('{}/{} model will be exported.'.format(index + 1, len(self.model_list)))
print('Model name is : ', mname)
print()
convert_deploy(model=mod,
backend_type=self.backend_type[deploy_backend],
dummy_input=self.dummy_input,
model_name=mname)