def to_torchserve(self,
filepath,
archive_format: tx.Literal["default",
"no-archive"] = "default"):
"""Build a TorchServe-compatible MAR file for this model."""
assert (
marmpu is not None
), "You must `pip install torch-model-archiver` to use this function."
os.makedirs(os.path.dirname(filepath), exist_ok=True)
with tempfile.TemporaryDirectory() as tdir:
serialized_file = os.path.join(tdir, "weights.pth")
index_to_name_file = os.path.join(tdir, "index_to_name.json")
model_file = os.path.join(tdir, "model.py")
torch.save(self.model.state_dict(prefix="model."), serialized_file)
with open(index_to_name_file, "w") as f:
f.write(json.dumps(self.serve_module_index))
with open(model_file, "w") as f:
f.write(self.serve_module_string)
args = types.SimpleNamespace(
model_name=os.path.basename(filepath),
serialized_file=serialized_file,
handler="object_detector",
model_file=model_file,
version="1.0",
requirements_file=None,
runtime="python",
extra_files=index_to_name_file,
export_path=os.path.dirname(filepath),
force=True,
archive_format=archive_format,
)
marmp.package_model(
args=args,
manifest=marmpu.ModelExportUtils.generate_manifest_json(args))
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)
