def save(model_in: MLModel):
"""Register a model into ModelDB and GridFS. `model.id` should be set as `None`, otherwise, the function will
raise a `ValueError`.
Args:
model_in (MLModelIn): model object to be registered
Return:
MLModel: Saved ML model object.
Raises:
BadRequestValueException: If `model.id` is not None.
ServiceException: If model has exists with the same primary keys (name, framework, engine and version).
"""
if _collection.count_documents(
filter=model_in.dict(
use_enum_values=True,
include={'architecture', 'framework', 'engine', 'version', 'task', 'dataset'}
),
limit=1
):
raise ServiceException(
f'Model with primary keys architecture={model_in.architecture}, '
f'framework={model_in.framework}, engine={model_in.engine}, version={model_in.version},'
f'task={model_in.task}, and dataset={model_in.dataset} has exists.'
)
# TODO: update weight ID in the MLModelIn
weight_id = _fs.put(bytes(model_in.weight), filename=model_in.weight.filename)
model = MLModel(**model_in.dict(exclude={'weight'}), weight=weight_id)
model.id = _collection.insert_one(model.dict(exclude_none=True, by_alias=True, use_enum_values=True)).inserted_id
return model
def _generate_model_family(
model: MLModel,
max_batch_size: int = -1
):
net = load(model.saved_path)
build_saved_dir_from_engine = partial(
generate_path_plain,
**model.dict(include={'architecture', 'framework', 'task', 'version'}),
)
inputs = model.inputs
outputs = model.outputs
model_input = model.model_input
generated_dir_list = list()
torchscript_dir = build_saved_dir_from_engine(engine=Engine.TORCHSCRIPT)
tfs_dir = build_saved_dir_from_engine(engine=Engine.TFS)
onnx_dir = build_saved_dir_from_engine(engine=Engine.ONNX)
trt_dir = build_saved_dir_from_engine(engine=Engine.TRT)
if isinstance(net, torch.nn.Module):
# to TorchScript
if converter.convert(net, 'pytorch', 'torchscript', save_path=torchscript_dir):
generated_dir_list.append(torchscript_dir.with_suffix('.zip'))
# to ONNX, TODO(lym): batch cache, input shape, opset version
if converter.convert(net, 'pytorch', 'onnx', save_path=onnx_dir, inputs=inputs, outputs=outputs, model_input=model_input, optimize=False):
generated_dir_list.append(onnx_dir.with_suffix('.onnx'))
# to TRT
# TRTConverter.from_onnx(
# onnx_path=onnx_dir.with_suffix('.onnx'), save_path=trt_dir, inputs=inputs, outputs=outputs
# )
elif isinstance(net, tf.keras.Model):
# to TFS
converter.convert(net, 'tensorflow', 'tfs', save_path=tfs_dir)
generated_dir_list.append(tfs_dir.with_suffix('.zip'))
# to TRT
converter.convert(net, 'tfs', 'trt', tf_path=tfs_dir, save_path=trt_dir, inputs=inputs, outputs=outputs, max_batch_size=32)
generated_dir_list.append(trt_dir.with_suffix('.zip'))
return generated_dir_list
def generate_model_family(model: MLModel, max_batch_size: int = -1):
model_weight_path = model.saved_path
if not Path(model.saved_path).exists():
(filepath, filename) = os.path.split(model.saved_path)
os.makedirs(filepath)
with open(model.saved_path, 'wb') as f:
f.write(model.weight.__bytes__())
net = load(model_weight_path)
build_saved_dir_from_engine = partial(
generate_path_plain,
**model.dict(include={'architecture', 'framework', 'task', 'version'}),
)
inputs = model.inputs
outputs = model.outputs
model_input = model.model_input
generated_dir_list = list()
torchscript_dir = build_saved_dir_from_engine(engine=Engine.TORCHSCRIPT)
tfs_dir = build_saved_dir_from_engine(engine=Engine.TFS)
onnx_dir = build_saved_dir_from_engine(engine=Engine.ONNX)
trt_dir = build_saved_dir_from_engine(engine=Engine.TRT)
if isinstance(net, torch.nn.Module):
_torchfamily(net, False, torchscript_dir, onnx_dir, generated_dir_list,
inputs, outputs, model_input)
elif isinstance(net, tf.keras.Model):
_tffamily(net, tfs_dir, generated_dir_list, trt_dir, inputs, outputs)
elif isinstance(net, xgb.XGBModel):
_xgbfamily(net, inputs, onnx_dir, generated_dir_list, torchscript_dir,
outputs, model_input)
elif isinstance(net, lgb.LGBMModel):
_lgbfamily(net, inputs, onnx_dir, generated_dir_list, torchscript_dir,
outputs, model_input)
elif isinstance(net, skl.base.BaseEstimator):
_sklfamily(net, inputs, onnx_dir, generated_dir_list, torchscript_dir,
outputs, model_input)
return generated_dir_list
def register_model(
model: MLModel,
convert: bool = True,
profile: bool = True,
) -> List[MLModel]:
"""Upload a model to ModelDB.
This function will upload the given model into the database with some variation. It may optionally generate a
branch of models (i.e. model family) with different optimization techniques. Besides, a benchmark will be
scheduled for each generated model, in order to gain profiling results for model selection strategies.
In the `no_generate` model(i.e. `no_generate` flag is set to be `True`), `architecture`, `framework`, `engine`
and `version` could be None. If any of the above arguments is `None`, all of them will be auto induced
from the origin_model path. An `ValueError` will be raised if the mata info cannot be induced.
TODO:
This function has a super comprehensive logic, need to be simplified.
Arguments:
model: Required inputs for register a model. All information is wrapped in such model.
convert (bool): Flag for generation of model family. Default to True.
profile (bool): Flag for profiling uploaded (including converted) models. Default to True.
"""
models = list()
model_dir_list = list()
model.model_status = [ModelStatus.PUBLISHED]
models.append(save(model))
# generate model family
if convert:
model_dir_list.extend(converter.generate_model_family(model))
# register
model_data = model.dict(exclude={'weight', 'id', 'model_status', 'engine'})
for model_dir in model_dir_list:
parse_result = parse_path_plain(model_dir)
engine = parse_result['engine']
model_cvt = MLModel(**model_data,
weight=model_dir,
engine=engine,
model_status=[ModelStatus.CONVERTED])
models.append(save(model_cvt))
# profile registered model
if profile:
from modelci.controller import job_executor
from modelci.controller.executor import Job
file = tf.keras.utils.get_file(
"grace_hopper.jpg",
"https://storage.googleapis.com/download.tensorflow.org/example_images/grace_hopper.jpg"
)
test_img_bytes = cv2.imread(file)
kwargs = {
'repeat_data': test_img_bytes,
'batch_size': 32,
'batch_num': 100,
'asynchronous': False,
}
for model in models:
model.model_status = [ModelStatus.PROFILING]
ModelService.update_model(model)
kwargs['model_info'] = model
engine = model.engine
if engine == Engine.TORCHSCRIPT:
client = CVTorchClient(**kwargs)
elif engine == Engine.TFS:
client = CVTFSClient(**kwargs)
elif engine == Engine.ONNX:
client = CVONNXClient(**kwargs)
elif engine == Engine.TRT:
client = CVTRTClient(**kwargs)
else:
raise ValueError(f'No such serving engine: {engine}')
job_cuda = Job(client=client, device='cuda:0', model_info=model)
# job_cpu = Job(client=client, device='cpu', model_info=model)
job_executor.submit(job_cuda)
# job_executor.submit(job_cpu)
return models