def test_manifest_json(self): manifest = ModelExportUtils.generate_manifest_json(self.args) manifest_json = json.loads(manifest) assert manifest_json['runtime'] == RuntimeType.PYTHON.value assert 'engine' in manifest_json assert 'model' in manifest_json assert 'publisher' in manifest_json assert 'license' not in manifest_json
def test_export_model_method(self, patches): patches.export_utils.check_mar_already_exists.return_value = '/Users/dummyUser/' patches.export_utils.check_custom_model_types.return_value = '/Users/dummyUser', ['a.txt', 'b.txt'] patches.export_utils.zip.return_value = None package_model(self.args, ModelExportUtils.generate_manifest_json(self.args)) patches.export_utils.validate_inputs.assert_called() patches.export_utils.archive.assert_called() patches.export_utils.clean_temp_files.assert_called()
def mmcls2torchserve( config_file: str, checkpoint_file: str, output_folder: str, model_name: str, model_version: str = '1.0', force: bool = False, ): """Converts mmclassification model (config + checkpoint) to TorchServe `.mar`. Args: config_file: In MMClassification config format. The contents vary for each task repository. checkpoint_file: In MMClassification 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. """ mmcv.mkdir_or_exist(output_folder) config = mmcv.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}/mmcls_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 mmgen2torchserver(config_file: str, checkpoint_file: str, output_folder: str, model_name: str, model_version: str = '1.0', model_type: str = 'unconditional', force: bool = False): """Converts MMGeneration model (config + checkpoint) to TorchServe `.mar`. Args: config_file (str): Path of config file. The config should in MMGeneration format. checkpoint_file (str): Path of checkpoint. The checkpoint should in MMGeneration checkpoint format. output_folder (str): Folder where `{model_name}.mar` will be created. The file created will be in TorchServe archive format. model_name (str): Name of the generated ``'mar'`` file. 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 (str, optional): Model's version. Defaults to '1.0'. model_type (str, optional): Type of the model to be convert. Handler named ``{model_type}_handler`` would be used to generate ``mar`` file. Defaults to 'unconditional'. force (bool, optional): If True, existing `{model_name}.mar` will be overwritten. Default to False. """ mmcv.mkdir_or_exist(output_folder) config = mmcv.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}/mmgen_{model_type}_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 pack(self, model, dst_path, model_name, model_version, additional_files=None): """ Package a model into the MAR archive so that it can be used for serving using TorchServing :param model: an object representing a model :type model: torch.nn.Module :param dst_path: the destination base path (do not include the filename) of the MAR :type dst_path: str :param model_name: the name of the model (will be also used to define the prediction endpoint) :type model_name: str :param model_version: a string encoding the version of the model :type model_version: str :param additional_files: an optional list of files that should be included in the MAR :type additional_files: iterable :return: None """ if additional_files is None: additional_files = [] # save model torch.save(model, os.path.join(self.tmp_dir, 'model.pt')) args = ArgClass( os.path.join(self.tmp_dir, 'model.pt'), self.handler, ','.join([ os.path.join(self.tmp_dir, 'pre_process_tform.pkl'), os.path.join(self.tmp_dir, 'post_process_tform.pkl'), os.path.join(self.tmp_dir, 'metadata.json') ] + additional_files), dst_path, model_name, model_version) manifest = ModelExportUtils.generate_manifest_json(args) package_model(args, manifest=manifest)
def mmpose2torchserve(config_file: str, checkpoint_file: str, output_folder: str, model_name: str, model_version: str = '1.0', force: bool = False): """Converts MMPose model (config + checkpoint) to TorchServe `.mar`. Args: config_file: In MMPose config format. The contents vary for each task repository. checkpoint_file: In MMPose 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. """ mmcv.mkdir_or_exist(output_folder) config = mmcv.Config.fromfile(config_file) with TemporaryDirectory() as tmpdir: model_file = osp.join(tmpdir, 'config.py') config.dump(model_file) handler_path = osp.join(osp.dirname(__file__), 'mmpose_handler.py') model_name = model_name or osp.splitext( osp.basename(checkpoint_file))[0] # use mmcv CheckpointLoader if checkpoint is not from a local file if not osp.isfile(checkpoint_file): ckpt = CheckpointLoader.load_checkpoint(checkpoint_file) checkpoint_file = osp.join(tmpdir, 'checkpoint.pth') with open(checkpoint_file, 'wb') as f: torch.save(ckpt, f) args = Namespace( **{ 'model_file': model_file, 'serialized_file': checkpoint_file, 'handler': handler_path, 'model_name': model_name, '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)