コード例 #1
0
ファイル: mxnet_sdk.py プロジェクト: lanking520/djlsdk
def export_mxnet(model: gluon.HybridBlock,
                 sample_input: List[mx.ndarray.NDArray]):
    model.hybridize()
    model(*sample_input)
    directory = "build"
    if os.path.exists(directory) and os.path.isdir(directory):
        shutil.rmtree(directory)
        os.mkdir(directory)
    model_path = os.path.join(directory, "mxmodel")
    # Create and stored the model
    model.export(model_path)
    zip_files(directory, model_path)
    # Start creating template
    array = [ndarray_to_numpy(x) for x in sample_input]
    dest = get_djl_template("mxmodel.zip", array, add_deps())
    shutil.copy(model_path + ".zip", dest)
コード例 #2
0
ファイル: gluon.py プロジェクト: aarnphm/BentoML
def save(
    name: str,
    model: gluon.HybridBlock,
    *,
    labels: t.Optional[t.Dict[str, str]] = None,
    custom_objects: t.Optional[t.Dict[str, t.Any]] = None,
    metadata: t.Optional[t.Dict[str, t.Any]] = None,
    model_store: "ModelStore" = Provide[BentoMLContainer.model_store],
) -> Tag:
    """
    Save a model instance to BentoML modelstore.

    Args:
        name (:code:`str`):
            Name for given model instance. This should pass Python identifier check.
        model (`mxnet.gluon.HybridBlock`):
            Instance of gluon.HybridBlock model to be saved.
        labels (:code:`Dict[str, str]`, `optional`, default to :code:`None`):
            user-defined labels for managing models, e.g. team=nlp, stage=dev
        custom_objects (:code:`Dict[str, Any]]`, `optional`, default to :code:`None`):
            user-defined additional python objects to be saved alongside the model,
            e.g. a tokenizer instance, preprocessor function, model configuration json
        metadata (:code:`Dict[str, Any]`, `optional`,  default to :code:`None`):
            Custom metadata for given model.
        model_store (:mod:`~bentoml._internal.models.store.ModelStore`, default to :mod:`BentoMLContainer.model_store`):
            BentoML modelstore, provided by DI Container.

    Returns:
        :obj:`~bentoml.Tag`: A :obj:`tag` with a format `name:version` where `name` is the user-defined model's name, and a generated `version` by BentoML.

    Examples:

    .. code-block:: python

        import mxnet
        import mxnet.gluon as gluon
        import bentoml


        def train_gluon_classifier() -> gluon.nn.HybridSequential:
            net = mxnet.gluon.nn.HybridSequential()
            net.hybridize()
            net.forward(mxnet.nd.array(0))
            return net

        model = train_gluon_classifier()

        tag = bentoml.gluon.save("gluon_block", model)
    """  # noqa

    context: t.Dict[str, t.Any] = {
        "framework_name": "gluon",
        "pip_dependencies": [f"mxnet=={get_pkg_version('mxnet')}"],
    }
    options: t.Dict[str, t.Any] = dict()
    with bentoml.models.create(
            name,
            module=MODULE_NAME,
            labels=labels,
            custom_objects=custom_objects,
            options=options,
            context=context,
            metadata=metadata,
    ) as _model:

        model.export(_model.path_of(SAVE_NAMESPACE))

        return _model.tag