mlrun.frameworks.tf_keras#

mlrun.frameworks.tf_keras.apply_mlrun(model: tensorflow.keras.Model | None = None, model_name: str | None = None, tag: str = '', model_path: str | None = None, model_format: str = 'SavedModel', save_traces: bool = False, modules_map: dict[str, Union[NoneType, str, list[str]]] | str | None = None, custom_objects_map: dict[str, Union[str, list[str]]] | str | None = None, custom_objects_directory: str | None = None, context: MLClientCtx | None = None, auto_log: bool = True, tensorboard_directory: str | None = None, mlrun_callback_kwargs: dict[str, Any] | None = None, tensorboard_callback_kwargs: dict[str, Any] | None = None, use_horovod: bool | None = None, **kwargs) → TFKerasModelHandler[source]#

Wrap the given model with MLRun's interface providing it with mlrun's additional features.

Parameters:

  • model -- The model to wrap. Can be loaded from the model path given as well.

  • model_name -- The model name to use for storing the model artifact. If not given, the tf.keras.Model.name will be used.

  • tag -- The model's tag to log with.

  • model_path -- The model's store object path. Mandatory for evaluation (to know which model to update). If model is not provided, it will be loaded from this path.

  • model_format -- The format to use for saving and loading the model. Should be passed as a member of the class 'ModelFormats'. Default: 'ModelFormats.SAVED_MODEL'.

  • save_traces -- Whether or not to use functions saving (only available for the 'SavedModel' format) for loading the model later without the custom objects dictionary. Only from tensorflow version >= 2.4.0. Using this setting will increase the model saving size.

  • modules_map --

    A dictionary of all the modules required for loading the model. Each key is a path to a module and its value is the object name to import from it. All the modules will be imported globally. If multiple objects needed to be imported from the same module a list can be given. The map can be passed as a path to a json file as well. For example:

    {
    "module1": None,  # import module1
    "module2": ["func1", "func2"],  # from module2 import func1, func2
    "module3.sub_module": "func3",  # from module3.sub_module import func3
}

    If the model path given is of a store object, the modules map will be read from the logged modules map artifact of the model.

  • custom_objects_map --

    A dictionary of all the custom objects required for loading the model. Each key is a path to a python file and its value is the custom object name to import from it. If multiple objects needed to be imported from the same py file a list can be given. The map can be passed as a path to a json file as well. For example:

    {
    "/.../custom_optimizer.py": "optimizer",
    "/.../custom_layers.py": ["layer1", "layer2"],
}

    All the paths will be accessed from the given 'custom_objects_directory', meaning each py file will be read from 'custom_objects_directory/<MAP VALUE>'. If the model path given is of a store object, the custom objects map will be read from the logged custom object map artifact of the model. Notice: The custom objects will be imported in the order they came in this dictionary (or json). If a custom object is depended on another, make sure to put it below the one it relies on.

  • custom_objects_directory -- Path to the directory with all the python files required for the custom objects. Can be passed as a zip file as well (will be extracted during the run before loading the model). If the model path given is of a store object, the custom objects files will be read from the logged custom object artifact of the model.

  • context -- MLRun context to work with. If no context is given it will be retrieved via 'mlrun.get_or_create_ctx(None)'

  • auto_log -- Whether or not to apply MLRun's auto logging on the model. Default: True.

  • tensorboard_directory -- If context is not given, or if wished to set the directory even with context, this will be the output for the event logs of tensorboard. If not given, the 'tensorboard_dir' parameter will be tried to be taken from the provided context. If not found in the context, the default tensorboard output directory will be: /User/.tensorboard/<PROJECT_NAME> or if working on local, the set artifacts path.

  • mlrun_callback_kwargs -- Key word arguments for the MLRun callback. For further information see the documentation of the class 'MLRunLoggingCallback'. Note that both 'context' and 'auto_log' parameters are already given here.

  • tensorboard_callback_kwargs -- Key word arguments for the tensorboard callback. For further information see the documentation of the class 'TensorboardLoggingCallback'. Note that both 'context' and 'auto_log' parameters are already given here.

  • use_horovod -- Whether or not to use horovod - a distributed training framework. Default: None, meaning it will be read from context if available and if not - False.

Returns:

The model with MLRun's interface.