def _pack_obj(obj): """Recursively packs objects. """ try: return copy.deepcopy(obj) except TypeError: pass # is this a Keras serializable? try: model_metadata = saving_utils.model_metadata(obj) training_config = model_metadata["training_config"] model = serialize(obj) weights = obj.get_weights() return SavedKerasModel( cls=obj.__class__, model=model, weights=weights, training_config=training_config, ) except (TypeError, AttributeError): pass # try manually packing the object if hasattr(obj, "__dict__"): for key, val in obj.__dict__.items(): obj.__dict__[key] = _pack_obj(val) return obj if isinstance(obj, (list, tuple)): obj_type = type(obj) new_obj = obj_type([_pack_obj(o) for o in obj]) return new_obj return obj
def save_optimizer_weights(model, filepath, overwrite=True, **kwargs): if not isinstance(filepath, h5py.File): # If file exists and should not be overwritten. if not overwrite and os.path.isfile(filepath): proceed = hdf5_format.ask_to_proceed_with_overwrite(filepath) if not proceed: return f = h5py.File(filepath, mode='w') opened_new_file = True else: f = filepath opened_new_file = False try: model_metadata = saving_utils.model_metadata( model, include_optimizer=True, require_config=False) for k, v in model_metadata.items(): if isinstance(v, (dict, list, tuple)): f.attrs[k] = json.dumps( v, default=serialization.get_json_type).encode('utf8') else: f.attrs[k] = v if not isinstance(model.optimizer, optimizers.TFOptimizer): hdf5_format.save_optimizer_weights_to_hdf5_group(f, model.optimizer) f.flush() finally: if opened_new_file: f.close()
def _python_properties_internal(self): metadata = super(ModelSavedModelSaver, self)._python_properties_internal() metadata.update( saving_utils.model_metadata(self.obj, include_optimizer=True, require_config=False)) return metadata
def _write_array( self, include_optimizer: bool, serialized_weights: bytes, serialized_optimizer_weights: bytes, meta: Optional[Meta], ) -> None: """Write Tensorflow model to a TileDB array.""" assert self.model # TODO: Change timestamp when issue in core is resolved with tiledb.open(self.uri, "w", timestamp=current_milli_time(), ctx=self.ctx) as tf_model_tiledb: if isinstance(self.model, (Functional, Sequential)): tf_model_tiledb[:] = { "model_weights": np.array([serialized_weights]), "optimizer_weights": np.array([serialized_optimizer_weights]), } else: # Insert weights and optimizer layer_names = [] weight_names = [] weight_values = [] for layer in sorted(self.model.layers, key=attrgetter("name")): weights = layer.trainable_weights + layer.non_trainable_weights weight_values.append( pickle.dumps(backend.batch_get_value(weights))) weight_names.append( pickle.dumps([w.name.encode("utf8") for w in weights])) layer_names.append(layer.name) tf_model_tiledb[:] = { "layer_name": np.array(layer_names), "weight_values": np.array(weight_values), "weight_names": np.array(weight_names), # TODO (TeamML) Fix array scheme to avoid optimizer_weight repetition. Nullable "optimizer_weights": np.array([ serialized_optimizer_weights for _ in range(len(self.model.layers)) ]), } # Insert all model metadata model_metadata = saving_utils.model_metadata( self.model, include_optimizer) for key, value in model_metadata.items(): tf_model_tiledb.meta[key] = json.dumps( value, default=json_utils.get_json_type).encode("utf8") self.update_model_metadata(array=tf_model_tiledb, meta=meta)
def _python_properties_internal(self): metadata = super(ModelSavedModelSaver, self)._python_properties_internal() # Network stateful property is dependent on the child layers. metadata.pop('stateful') metadata['is_graph_network'] = self.obj._is_graph_network # pylint: disable=protected-access metadata.update( saving_utils.model_metadata(self.obj, include_optimizer=True, require_config=False)) return metadata
def pack_keras_model(model_obj, protocol): """Pickle a Keras Model. Arguments: model_obj: an instance of a Keras Model. protocol: pickle protocol version, ignored. Returns ------- Pickled model A tuple following the pickle protocol. """ model_metadata = saving_utils.model_metadata(model_obj) training_config = model_metadata.get("training_config", None) model = serialize_layer(model_obj) weights = model_obj.get_weights() return (unpack_keras_model, (model, training_config, weights))
def _write_array( self, model: Model, include_optimizer: bool, serialized_weights: bytes, serialized_optimizer_weights: bytes, meta: Optional[dict], ): """ Writes Tensorflow model to a TileDB array. """ with tiledb.open(self.uri, "w") as tf_model_tiledb: # Insert weights and optimizer tf_model_tiledb[:] = { "model_weights": np.array([serialized_weights]), "optimizer_weights": np.array([serialized_optimizer_weights]), } # Insert all model metadata model_metadata = saving_utils.model_metadata( model, include_optimizer) for key, value in model_metadata.items(): try: tf_model_tiledb.meta[key] = json.dumps( value, default=json_utils.get_json_type).encode("utf8") except: logging.warning( "Exception occurred during Json serialization of metadata!" ) # Add Python version to array's metadata tf_model_tiledb.meta["python_version"] = platform.python_version() # Add extra metadata given by the user to array's metadata if meta: for key, value in meta.items(): try: tf_model_tiledb.meta[key] = json.dumps( value, default=json_utils.get_json_type).encode("utf8") except: logging.warning( "Exception occurred during Json serialization of metadata!" )
def pack_keras_model(model_obj, protocol): """Pickle a Keras Model. Arguments: model_obj: an instance of a Keras Model. protocol: pickle protocol version, ignored. Returns ------- Pickled model A tuple following the pickle protocol. """ if not isinstance(model_obj, Model): raise TypeError("`model_obj` must be an instance of a Keras Model") # pack up model model_metadata = saving_utils.model_metadata(model_obj) training_config = model_metadata.get("training_config", None) model = serialize(model_obj) weights = model_obj.get_weights() return (unpack_keras_model, (model, training_config, weights))
def _clone_prebuilt_model(build_fn): """Clones and compiles a pre-built model when build_fn is an existing Keras model instance. Arguments: build_fn : instance of Keras Model. Returns: copy of the input model with no training. """ model = clone_model(build_fn) # clone_model does not compy over compilation parameters, do those manually model_metadata = saving_utils.model_metadata(build_fn) if "training_config" in model_metadata: training_config = model_metadata["training_config"] else: raise ValueError("To use %s as `build_fn`, you must compile" "it first." % build_fn) model.compile( **saving_utils.compile_args_from_training_config(training_config)) return model
def save_model_to_hdf5(model, filepath, overwrite=True, include_optimizer=True): """Saves a model to a HDF5 file. The saved model contains: - the model's configuration (topology) - the model's weights - the model's optimizer's state (if any) Thus the saved model can be reinstantiated in the exact same state, without any of the code used for model definition or training. Arguments: model: Keras model instance to be saved. filepath: One of the following: - String, path where to save the model - `h5py.File` object where to save the model overwrite: Whether we should overwrite any existing model at the target location, or instead ask the user with a manual prompt. include_optimizer: If True, save optimizer's state together. Raises: ImportError: if h5py is not available. """ if h5py is None: raise ImportError('`save_model` requires h5py.') # TODO(psv) Add warning when we save models that contain non-serializable # entities like metrics added using `add_metric` and losses added using # `add_loss.` if len(model.weights) != len(model._undeduplicated_weights): logging.warning( 'Found duplicated `Variable`s in Model\'s `weights`. ' 'This is usually caused by `Variable`s being shared by ' 'Layers in the Model. These `Variable`s will be treated ' 'as separate `Variable`s when the Model is restored. To ' 'avoid this, please save with `save_format="tf"`.') if not isinstance(filepath, h5py.File): # If file exists and should not be overwritten. if not overwrite and os.path.isfile(filepath): proceed = ask_to_proceed_with_overwrite(filepath) if not proceed: return f = h5py.File(filepath, mode='w') opened_new_file = True else: f = filepath opened_new_file = False try: model_metadata = saving_utils.model_metadata(model, include_optimizer) for k, v in model_metadata.items(): if isinstance(v, (dict, list, tuple)): f.attrs[k] = json.dumps( v, default=serialization.get_json_type).encode('utf8') else: f.attrs[k] = v model_weights_group = f.create_group('model_weights') model_layers = model.layers save_weights_to_hdf5_group(model_weights_group, model_layers) # TODO(b/128683857): Add integration tests between tf.keras and external # Keras, to avoid breaking TF.js users. if (include_optimizer and model.optimizer and not isinstance(model.optimizer, optimizers.TFOptimizer)): save_optimizer_weights_to_hdf5_group(f, model.optimizer) f.flush() finally: if opened_new_file: f.close()
def __reduce__(self): model_metadata = saving_utils.model_metadata(self) training_config = model_metadata.get("training_config", None) model = serialize(self) weights = self.get_weights() return (unpack, (model, training_config, weights))
def _reduce_tf_model(model): model_metadata = saving_utils.model_metadata(model) training_config = model_metadata.get("training_config", None) weights = model.get_weights() model = serialize(model) return TensorflowDispatcher._make_model, (model, training_config, weights)