| | |
| | from argparse import ArgumentParser, Namespace |
| | from pathlib import Path |
| | from tempfile import TemporaryDirectory |
| |
|
| | from mmengine.config import Config |
| | from mmengine.utils import mkdir_or_exist |
| |
|
| | try: |
| | from model_archiver.model_packaging import package_model |
| | from model_archiver.model_packaging_utils import ModelExportUtils |
| | except ImportError: |
| | package_model = None |
| |
|
| |
|
| | def mmdet2torchserve( |
| | config_file: str, |
| | checkpoint_file: str, |
| | output_folder: str, |
| | model_name: str, |
| | model_version: str = '1.0', |
| | force: bool = False, |
| | ): |
| | """Converts MMDetection model (config + checkpoint) to TorchServe `.mar`. |
| | |
| | Args: |
| | config_file: |
| | In MMDetection config format. |
| | The contents vary for each task repository. |
| | checkpoint_file: |
| | In MMDetection checkpoint format. |
| | The contents vary for each task repository. |
| | output_folder: |
| | Folder where `{model_name}.mar` will be created. |
| | The file created will be in TorchServe archive format. |
| | model_name: |
| | If not None, used for naming the `{model_name}.mar` file |
| | that will be created under `output_folder`. |
| | If None, `{Path(checkpoint_file).stem}` will be used. |
| | model_version: |
| | Model's version. |
| | force: |
| | If True, if there is an existing `{model_name}.mar` |
| | file under `output_folder` it will be overwritten. |
| | """ |
| | mkdir_or_exist(output_folder) |
| |
|
| | config = Config.fromfile(config_file) |
| |
|
| | with TemporaryDirectory() as tmpdir: |
| | config.dump(f'{tmpdir}/config.py') |
| |
|
| | args = Namespace( |
| | **{ |
| | 'model_file': f'{tmpdir}/config.py', |
| | 'serialized_file': checkpoint_file, |
| | 'handler': f'{Path(__file__).parent}/mmdet_handler.py', |
| | 'model_name': model_name or Path(checkpoint_file).stem, |
| | 'version': model_version, |
| | 'export_path': output_folder, |
| | 'force': force, |
| | 'requirements_file': None, |
| | 'extra_files': None, |
| | 'runtime': 'python', |
| | 'archive_format': 'default' |
| | }) |
| | manifest = ModelExportUtils.generate_manifest_json(args) |
| | package_model(args, manifest) |
| |
|
| |
|
| | def parse_args(): |
| | parser = ArgumentParser( |
| | description='Convert MMDetection models to TorchServe `.mar` format.') |
| | parser.add_argument('config', type=str, help='config file path') |
| | parser.add_argument('checkpoint', type=str, help='checkpoint file path') |
| | parser.add_argument( |
| | '--output-folder', |
| | type=str, |
| | required=True, |
| | help='Folder where `{model_name}.mar` will be created.') |
| | parser.add_argument( |
| | '--model-name', |
| | type=str, |
| | default=None, |
| | help='If not None, used for naming the `{model_name}.mar`' |
| | 'file that will be created under `output_folder`.' |
| | 'If None, `{Path(checkpoint_file).stem}` will be used.') |
| | parser.add_argument( |
| | '--model-version', |
| | type=str, |
| | default='1.0', |
| | help='Number used for versioning.') |
| | parser.add_argument( |
| | '-f', |
| | '--force', |
| | action='store_true', |
| | help='overwrite the existing `{model_name}.mar`') |
| | args = parser.parse_args() |
| |
|
| | return args |
| |
|
| |
|
| | if __name__ == '__main__': |
| | args = parse_args() |
| |
|
| | if package_model is None: |
| | raise ImportError('`torch-model-archiver` is required.' |
| | 'Try: pip install torch-model-archiver') |
| |
|
| | mmdet2torchserve(args.config, args.checkpoint, args.output_folder, |
| | args.model_name, args.model_version, args.force) |
| |
|
