mlrun.frameworks.tf_keras#
- mlrun.frameworks.tf_keras.apply_mlrun(model: Optional[tensorflow.keras.Model] = None, model_name: Optional[str] = None, tag: str = '', model_path: Optional[str] = None, model_format: str = 'SavedModel', save_traces: bool = False, modules_map: Optional[Union[Dict[str, Union[None, str, List[str]]], str]] = None, custom_objects_map: Optional[Union[Dict[str, Union[str, List[str]]], str]] = None, custom_objects_directory: Optional[str] = None, context: Optional[MLClientCtx] = None, auto_log: bool = True, tensorboard_directory: Optional[str] = None, mlrun_callback_kwargs: Optional[Dict[str, Any]] = None, tensorboard_callback_kwargs: Optional[Dict[str, Any]] = None, use_horovod: Optional[bool] = 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.