def _strip_clustering_wrapper(layer): if isinstance(layer, cluster_wrapper.ClusterWeights): if not hasattr(layer.layer, '_batch_input_shape') and\ hasattr(layer, '_batch_input_shape'): layer.layer._batch_input_shape = layer._batch_input_shape # We reset both arrays of weights, so that we can guarantee the correct # order of newly created weights layer.layer._trainable_weights = [] layer.layer._non_trainable_weights = [] for i in range(len(layer.restore)): # This is why we used integers as keys name, weight = layer.restore[i] # In both cases we use k.batch_get_value since we need physical copies # of the arrays to initialize a new tensor if i in layer.gone_variables: # If the variable was removed because it was clustered, we restore it # by using updater we created earlier new_weight_value = k.batch_get_value([weight()])[0] else: # If the value was not clustered(e.g. bias), we still store a valid # reference to the tensor. We use this reference to get the value new_weight_value = k.batch_get_value([weight])[0] layer.layer.add_weight( name=name, shape=new_weight_value.shape, initializer=initializers.Constant(new_weight_value), trainable=True) # When all weights are filled with the values, just return the underlying # layer since it is now fully autonomous from its wrapper return layer.layer return layer
def F(inputs): self.count += 1 R = fast_train_function(inputs) if self.count % self.k == 0: K.batch_get_value(slow_updates) K.batch_get_value(copy_updates) return R
def on_train_batch_end(self, batch, logs=None): self.count += 1 if self.slow_weights is None: with tf.control_dependencies(self.model.trainable_weights): self.slow_weights = [] for fast_param in self.model.trainable_weights: with ops.control_dependencies([fast_param]): slow_param = tf.Variable(fast_param.initialized_value(), dtype=fast_param.dtype, trainable=False, name=fast_param.name.split(":")[0]) self.slow_weights.append(slow_param) K.track_variable(slow_param) else: if self.count % self.k == 0: slow_ups, fast_ups = [], [] for fast, slow in zip(self.model.trainable_weights, self.slow_weights): slow_ups.append(K.update(slow, slow + self.alpha * (fast - slow))) with tf.control_dependencies(slow_ups): for fast, slow in zip(self.model.trainable_weights, self.slow_weights): fast_ups.append(K.update(fast, slow)) K.batch_get_value(slow_ups) K.batch_get_value(fast_ups)
def _assert_weights_equal_value(self, annotated_weights, emulated_weights): annotated_weight_values = K.batch_get_value(annotated_weights) emulated_weight_values = K.batch_get_value(emulated_weights) self.assertEqual(len(annotated_weight_values), len(emulated_weight_values)) for aw, ew in zip(annotated_weight_values, emulated_weight_values): self.assertAllClose(aw, ew)
def test_save_load_with_dense_features(self, tmpdir, api, loss, optimizer, metrics): if optimizer is None: pytest.skip() cols = [ feature_column_lib.numeric_column("a"), feature_column_lib.indicator_column( feature_column_lib.categorical_column_with_vocabulary_list( "b", ["one", "two"])), ] input_layers = { "a": keras.layers.Input(shape=(1, ), name="a"), "b": keras.layers.Input(shape=(1, ), name="b", dtype="string"), } fc_layer = dense_features.DenseFeatures(cols)(input_layers) output = keras.layers.Dense(10)(fc_layer) model = keras.models.Model(input_layers, output) model.compile( loss=loss, optimizer=optimizer, metrics=[metrics], ) tiledb_uri = os.path.join(tmpdir, "model_array") tiledb_model_obj = TensorflowKerasTileDBModel(uri=tiledb_uri, model=model) tiledb_model_obj.save(include_optimizer=True) loaded_model = tiledb_model_obj.load(compile_model=True) model_opt_weights = batch_get_value(getattr(model.optimizer, "weights")) loaded_opt_weights = batch_get_value( getattr(loaded_model.optimizer, "weights")) # Assert optimizer weights are equal for weight_model, weight_loaded_model in zip(model_opt_weights, loaded_opt_weights): np.testing.assert_array_equal(weight_model, weight_loaded_model) inputs_a = np.arange(10).reshape(10, 1) inputs_b = np.arange(10).reshape(10, 1).astype("str") # Assert model predictions are equal np.testing.assert_array_equal( loaded_model.predict({ "a": inputs_a, "b": inputs_b }), model.predict({ "a": inputs_a, "b": inputs_b }), )
def test_get_and_set_weights_does_not_add_ops(self): table_handler = self.get_table_handler() table_data = {b'a': 1, b'b': 2, b'c': 3} table_handler.set_weights( [list(table_data.keys()), list(table_data.values())]) _ = backend.batch_get_value(table_handler.get_tensors()) backend.get_session().graph.finalize() table_handler.set_weights( [list(table_data.keys()), list(table_data.values())]) _ = backend.batch_get_value(table_handler.get_tensors())
def on_epoch_end(self, batch, logs=None): # At the end of every epoch, remask the weights. This ensures that when # the model is saved after completion, the weights represent mask*weights. layers = self.model.layers weight_mask_ops = [] for layer in layers: if isinstance(layer, pruning_wrapper.PruneLowMagnitude): if tf.executing_eagerly(): layer.pruning_obj.weight_mask_op() else: weight_mask_ops.append(layer.pruning_obj.weight_mask_op()) K.batch_get_value(weight_mask_ops)
def get_weights(self): """Returns the current weights of the optimizer. The weights of an optimizer are its state (ie, variables). This function returns the weight values associated with this optimizer as a list of Numpy arrays. The first value is always the iterations count of the optimizer, followed by the optimizer's state variables in the order they were created. The returned list can in turn be used to load state into similarly parameterized optimizers. For example, the RMSprop optimizer for this simple model returns a list of three values-- the iteration count, followed by the root-mean-square value of the kernel and bias of the single Dense layer: >>> opt = tf.keras.optimizers.RMSprop() >>> m = tf.keras.models.Sequential([tf.keras.layers.Dense(10)]) >>> m.compile(opt, loss='mse') >>> data = np.arange(100).reshape(5, 20) >>> labels = np.zeros(5) >>> print('Training'); results = m.fit(data, labels) Training ... >>> len(opt.get_weights()) 3 Returns: Weights values as a list of numpy arrays. """ params = self.weights return backend.batch_get_value(params)
def save_weights_to_hdf5_group(f, layers): """Saves the weights of a list of layers to a HDF5 group. Arguments: f: HDF5 group. layers: List of layer instances. """ from tensorflow.python.keras import __version__ as keras_version # pylint: disable=g-import-not-at-top save_attributes_to_hdf5_group( f, 'layer_names', [layer.name.encode('utf8') for layer in layers]) f.attrs['backend'] = K.backend().encode('utf8') f.attrs['keras_version'] = str(keras_version).encode('utf8') for layer in layers: g = f.create_group(layer.name) weights = _legacy_weights(layer) weight_values = K.batch_get_value(weights) weight_names = [w.name.encode('utf8') for w in weights] save_attributes_to_hdf5_group(g, 'weight_names', weight_names) for name, val in zip(weight_names, weight_values): param_dset = g.create_dataset(name, val.shape, dtype=val.dtype) if not val.shape: # scalar param_dset[()] = val else: param_dset[:] = val
def save_weights_to_hdf5_group(f, layers): """Saves the weights of a list of layers to a HDF5 group. Arguments: f: HDF5 group. layers: List of layer instances. """ from tensorflow.python.keras import __version__ as keras_version # pylint: disable=g-import-not-at-top save_attributes_to_hdf5_group( f, 'layer_names', [layer.name.encode('utf8') for layer in layers]) f.attrs['backend'] = K.backend().encode('utf8') f.attrs['keras_version'] = str(keras_version).encode('utf8') for layer in layers: g = f.create_group(layer.name) symbolic_weights = layer.weights weight_values = K.batch_get_value(symbolic_weights) weight_names = [] for i, (w, val) in enumerate(zip(symbolic_weights, weight_values)): if hasattr(w, 'name') and w.name: name = str(w.name) else: name = 'param_' + str(i) weight_names.append(name.encode('utf8')) save_attributes_to_hdf5_group(g, 'weight_names', weight_names) for name, val in zip(weight_names, weight_values): param_dset = g.create_dataset(name, val.shape, dtype=val.dtype) if not val.shape: # scalar param_dset[()] = val else: param_dset[:] = val
def set_weights(self, weights): """Sets the weights of the optimizer, from Numpy arrays. Should only be called after computing the gradients (otherwise the optimizer has no weights). Arguments: weights: a list of Numpy arrays. The number of arrays and their shape must match number of the dimensions of the weights of the optimizer (i.e. it should match the output of `get_weights`). Raises: ValueError: in case of incompatible weight shapes. """ params = self.weights if len(params) != len(weights): raise ValueError( 'Length of the specified weight list (' + str(len(weights)) + ') does not match the number of weights ' 'of the optimizer (' + str(len(params)) + ')') weight_value_tuples = [] param_values = K.batch_get_value(params) for pv, p, w in zip(param_values, params, weights): if pv.shape != w.shape: raise ValueError( 'Optimizer weight shape ' + str(pv.shape) + ' not compatible with ' 'provided weight shape ' + str(w.shape)) weight_value_tuples.append((p, w)) K.batch_set_value(weight_value_tuples)
def on_epoch_end(self, batch, logs=None): super(PruningSummaries, self).on_epoch_end(batch, logs) pruning_logs = {} params = [] layers = self.model.layers for layer in layers: if isinstance(layer, pruning_wrapper.PruneLowMagnitude): for _, mask, threshold in layer.pruning_vars: params.append(mask) params.append(threshold) params.append(self.model.optimizer.iterations) values = K.batch_get_value(params) iteration = values[-1] del values[-1] del params[-1] param_value_pairs = list(zip(params, values)) for mask, mask_value in param_value_pairs[::2]: pruning_logs.update( {mask.name + '/sparsity': 1 - np.mean(mask_value)}) for threshold, threshold_value in param_value_pairs[1::2]: pruning_logs.update( {threshold.name + '/threshold': threshold_value}) self._log_metrics(pruning_logs, '', iteration)
def set_weights(self, weights): """Sets the weights of the optimizer, from Numpy arrays. Should only be called after computing the gradients (otherwise the optimizer has no weights). Arguments: weights: a list of Numpy arrays. The number of arrays and their shape must match number of the dimensions of the weights of the optimizer (i.e. it should match the output of `get_weights`). Raises: ValueError: in case of incompatible weight shapes. """ params = self.weights if len(params) != len(weights): raise ValueError('Length of the specified weight list (' + str(len(weights)) + ') does not match the number of weights ' 'of the optimizer (' + str(len(params)) + ')') weight_value_tuples = [] param_values = K.batch_get_value(params) for pv, p, w in zip(param_values, params, weights): if pv.shape != w.shape: raise ValueError('Optimizer weight shape ' + str(pv.shape) + ' not compatible with ' 'provided weight shape ' + str(w.shape)) weight_value_tuples.append((p, w)) K.batch_set_value(weight_value_tuples)
def get_weights(self): """Returns the current value of the weights of the optimizer. Returns: A list of numpy arrays. """ return K.batch_get_value(self.weights)
def test_save_model_to_tiledb_array_weights(self, tmpdir, api, loss, optimizer, metrics): model = (api(num_hidden=1, num_classes=2, input_dim=3) if api != ConfigSubclassModel else api( hidden_units=[16, 16, 10])) tiledb_uri = os.path.join(tmpdir, "model_array") # Compiles the model if optimizer is present if optimizer: model.compile(loss=loss, optimizer=optimizer, metrics=[metrics]) input_shape = tuple(np.random.randint(20, size=2)) if not model.built: model.build(input_shape) tiledb_model_obj = TensorflowKerasTileDBModel(uri=tiledb_uri, model=model) tiledb_model_obj.save(include_optimizer=True if optimizer else False) loaded_model = ( tiledb_model_obj.load(compile_model=True if optimizer else False) if api != ConfigSubclassModel else tiledb_model_obj.load( compile_model=True if optimizer else False, custom_objects={"ConfigSubclassModel": ConfigSubclassModel}, input_shape=input_shape, )) data = np.random.rand( 100, input_shape[-1] if api == ConfigSubclassModel else 3) if optimizer: model_opt_weights = batch_get_value( getattr(model.optimizer, "weights")) loaded_opt_weights = batch_get_value( getattr(loaded_model.optimizer, "weights")) # Assert optimizer weights are equal for weight_model, weight_loaded_model in zip( model_opt_weights, loaded_opt_weights): np.testing.assert_array_equal(weight_model, weight_loaded_model) # Assert model predictions are equal np.testing.assert_array_equal(loaded_model.predict(data), model.predict(data))
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 get_weights(self): """Retrieves the weights of the model. Returns: A flat list of Numpy arrays. """ weights = [] for cell in self.cells: if isinstance(cell, Layer): weights += cell.weights return K.batch_get_value(weights)
def test_get_and_set_weights(self): table_handler = self.get_table_handler() table_data = {b'a': 1, b'b': 2, b'c': 3} table_handler.set_weights( [list(table_data.keys()), list(table_data.values())]) weights = backend.batch_get_value(table_handler.get_tensors()) weight_data = {key: value for key, value in zip(weights[0], weights[1])} self.assertDictEqual(table_data, weight_data)
def beam_search(predictions: List[np.ndarray], padding_token_id: int, start_sentence_token_id: int, end_sentence_token_id: int, beam_width: int = 5, beam_top_paths: int = 5): """ predictions: output from a softmax layer, y true labels # TODO if time permits implement own beam search, TF is too slow """ print("{}: In beam search".format(time.strftime(time.strftime("%Y-%m-%d-%H-%M%S")))) start_time = time.time() beam_search_predictions_list = [] beam_search_probs_list = [] for pred in predictions: top_path_prediction_tensors, probs = K.ctc_decode( np.expand_dims(pred, 0), (pred.shape[0],), greedy=False, beam_width=beam_width, top_paths=beam_top_paths ) beam_search_predictions_list.append(top_path_prediction_tensors) beam_search_probs_list.append(probs) # evaluate tensorflow graph print("{}: Evaluating beam search TF graph".format(time.strftime(time.strftime("%Y-%m-%d-%H-%M%S")))) beam_search_predictions_evaluated: List[np.ndarray] = K.batch_get_value(beam_search_predictions_list) print("{} Cleaning beamsearch results".format(time.strftime(time.strftime("%Y-%m-%d-%H-%M%S")))) best_predictions = [list(trim_pred(pred, padding_token_id, start_sentence_token_id, end_sentence_token_id) for pred in beam_search_single_result) for beam_search_single_result in beam_search_predictions_evaluated] del beam_search_predictions_evaluated # freeup much needed memory top_paths_predictions: np.ndarray = K.batch_get_value(beam_search_probs_list) best_predictions_probs = list(map(lambda pred: np.exp(pred[0]), top_paths_predictions)) del top_paths_predictions # freeup much needed memory print("beam search ended for one iteration in {}ms".format(time.time() - start_time)) return best_predictions, best_predictions_probs
def test_save_load_with_compile_functional(self): functional_model = testing_utils.get_small_functional_mlp( num_hidden=1, num_classes=2, input_dim=3) functional_model = add_optimizer(functional_model) tiledb_uri = os.path.join(self.get_temp_dir(), "model_array") tiledb_model_obj = TensorflowTileDB(uri=tiledb_uri) tiledb_model_obj.save(model=functional_model, include_optimizer=True) loaded_model = tiledb_model_obj.load(compile_model=True) data = np.random.rand(100, 3) model_opt_weights = batch_get_value( getattr(functional_model.optimizer, "weights")) loaded_opt_weights = batch_get_value( getattr(loaded_model.optimizer, "weights")) # Assert optimizer weights are equal for weight_model, weight_loaded_model in zip(model_opt_weights, loaded_opt_weights): np.testing.assert_array_equal(weight_model, weight_loaded_model) # Assert model predictions are equal np.testing.assert_array_equal(loaded_model.predict(data), functional_model.predict(data))
def set_weights(self, weights): """Set the weights of the optimizer. The weights of an optimizer are its state (ie, variables). This function takes the weight values associated with this optimizer as a list of Numpy arrays. The first value is always the iterations count of the optimizer, followed by the optimizer's state variables in the order they are created. The passed values are used to set the new state of the optimizer. For example, the RMSprop optimizer for this simple model takes a list of three values-- the iteration count, followed by the root-mean-square value of the kernel and bias of the single Dense layer: >>> opt = tf.keras.optimizers.RMSprop() >>> m = tf.keras.models.Sequential([tf.keras.layers.Dense(10)]) >>> m.compile(opt, loss='mse') >>> data = np.arange(100).reshape(5, 20) >>> labels = np.zeros(5) >>> print('Training'); results = m.fit(data, labels) Training ... >>> new_weights = [np.array(10), np.ones([20, 10]), np.zeros([10])] >>> opt.set_weights(new_weights) >>> opt.iterations <tf.Variable 'RMSprop/iter:0' shape=() dtype=int64, numpy=10> Arguments: weights: weight values as a list of numpy arrays. """ params = self.weights if len(params) != len(weights): raise ValueError( "You called `set_weights(weights)` on optimizer " + self._name + " with a weight list of length " + str(len(weights)) + ", but the optimizer was expecting " + str(len(params)) + " weights. Provided weights: " + str(weights)[:50] + "...") if not params: return weight_value_tuples = [] param_values = backend.batch_get_value(params) for pv, p, w in zip(param_values, params, weights): if pv.shape != w.shape: raise ValueError("Optimizer weight shape " + str(pv.shape) + " not compatible with " "provided weight shape " + str(w.shape)) weight_value_tuples.append((p, w)) backend.batch_set_value(weight_value_tuples)
def set_weights(self, weights): params = self.weights if len(params) != len(weights): raise ValueError( "You called `set_weights(weights)` on optimizer " + self._name + " with a weight list of length " + str(len(weights)) + ", but the optimizer was expecting " + str(len(params)) + " weights. Provided weights: " + str(weights)[:50] + "...") if not params: return weight_value_tuples = [] param_values = backend.batch_get_value(params) for pv, p, w in zip(param_values, params, weights): if pv.shape != w.shape: raise ValueError("Optimizer weight shape " + str(pv.shape) + " not compatible with " "provided weight shape " + str(w.shape)) weight_value_tuples.append((p, w)) backend.batch_set_value(weight_value_tuples)
def save_weights_to_hdf5_group(f, layers): """Saves the weights of a list of layers to a HDF5 group. Arguments: f: HDF5 group. layers: List of layer instances. """ from tensorflow.python.keras import __version__ as keras_version # pylint: disable=g-import-not-at-top save_attributes_to_hdf5_group( f, 'layer_names', [layer.name.encode('utf8') for layer in layers]) f.attrs['backend'] = K.backend().encode('utf8') f.attrs['keras_version'] = str(keras_version).encode('utf8') # On TPUs, modifying the graph between session.runs() triggers some expensive # recompilation overhead. To avoid this, we build up the full set of tensors # to save before fetching weights, thus only modifying the graph once. layer_weights_dict = {} for layer in layers: layer_weights_dict[layer.name] = [ ops.convert_to_tensor(w) for w in layer.weights ] for layer in layers: g = f.create_group(layer.name) symbolic_weights = layer_weights_dict[layer.name] weight_values = K.batch_get_value(symbolic_weights) weight_names = [] for i, (w, val) in enumerate(zip(symbolic_weights, weight_values)): if hasattr(w, 'name') and w.name: name = str(w.name) else: name = 'param_' + str(i) weight_names.append(name.encode('utf8')) save_attributes_to_hdf5_group(g, 'weight_names', weight_names) for name, val in zip(weight_names, weight_values): param_dset = g.create_dataset(name, val.shape, dtype=val.dtype) if not val.shape: # scalar param_dset[()] = val else: param_dset[:] = val
def train(self) -> Dict: for data_set in self.data_sets: data_set('train') valid_exists = all( data_set.valid_X is not None and data_set.valid_Y is not None for data_set in self.data_sets) self._model.fit( **{ 'x': self.data_sets[0].data_X, 'y': self.data_sets[0].data_Y, 'batch_size': self.batch_size, 'validation_data': ( self.data_sets[0].valid_X, self.data_sets[0].valid_Y) if valid_exists else None, 'epochs': self.epochs, 'verbose': 0, 'callbacks': [ tf.keras.callbacks.ModelCheckpoint(save_weights_only=True, save_best_only=True, filepath=self.tmp_file, verbose=0), tf.keras.callbacks. EarlyStopping(monitor='val_loss', patience=10) ] if valid_exists else None, }) if valid_exists: self._model.load_weights(self.tmp_file) state = dict() for f, (tfObj, vname_dict) in self._id2tfObj.items(): variables = tfObj.variables state[f] = { vname_dict[variable.name]: value for variable, value in zip(variables, K.batch_get_value(variables)) } state[NeuralNetworkFrameworkInterface.arg_CMP] = self.cmp return state
def save_tf_keras_optimizer(optimizer, h5py_file): if isinstance(optimizer, optimizers.TFOptimizer): logging.warning( 'TensorFlow optimizers do not ' 'make it possible to access ' 'optimizer attributes or optimizer state ' 'after instantiation. ' 'As a result, we cannot save the optimizer ' 'as part of the model save file.' 'You will have to compile your model again after loading it. ' 'Prefer using a Keras optimizer instead ' '(see keras.io/optimizers).') else: h5py_file.attrs['training_config'] = json.dumps( { 'optimizer_config': { 'class_name': optimizer.__class__.__name__, 'config': optimizer.get_config() } }, default=serialization.get_json_type).encode('utf8') # Save optimizer weights. symbolic_weights = getattr(optimizer, 'weights') if symbolic_weights: optimizer_weights_group = h5py_file.create_group( 'optimizer_weights') weight_values = K.batch_get_value(symbolic_weights) weight_names = [] for w, val in zip(symbolic_weights, weight_values): name = str(w.name) weight_names.append(name.encode('utf8')) optimizer_weights_group.attrs['weight_names'] = weight_names for name, val in zip(weight_names, weight_values): param_dset = optimizer_weights_group.create_dataset( name, val.shape, dtype=val.dtype) if not val.shape: # scalar param_dset[()] = val else: param_dset[:] = val h5py_file.flush()
def save_weights_to_hdf5_group(f, layers): """Saves the weights of a list of layers to a HDF5 group. Arguments: f: HDF5 group. layers: List of layer instances. """ from tensorflow.python.keras import __version__ as keras_version # pylint: disable=g-import-not-at-top save_attributes_to_hdf5_group( f, 'layer_names', [layer.name.encode('utf8') for layer in layers]) f.attrs['backend'] = K.backend().encode('utf8') f.attrs['keras_version'] = str(keras_version).encode('utf8') # On TPUs, modifying the graph between session.runs() triggers some expensive # recompilation overhead. To avoid this, we build up the full set of tensors # to save before fetching weights, thus only modifying the graph once. layer_weights_dict = {} for layer in layers: layer_weights_dict[layer.name] = [ops.convert_to_tensor(w) for w in layer.weights] for layer in layers: g = f.create_group(layer.name) symbolic_weights = layer_weights_dict[layer.name] weight_values = K.batch_get_value(symbolic_weights) weight_names = [] for i, (w, val) in enumerate(zip(symbolic_weights, weight_values)): if hasattr(w, 'name') and w.name: name = str(w.name) else: name = 'param_' + str(i) weight_names.append(name.encode('utf8')) save_attributes_to_hdf5_group(g, 'weight_names', weight_names) for name, val in zip(weight_names, weight_values): param_dset = g.create_dataset(name, val.shape, dtype=val.dtype) if not val.shape: # scalar param_dset[()] = val else: param_dset[:] = val
def save_optimizer_weights_to_hdf5_group(hdf5_group, optimizer): """Saves optimizer weights of a optimizer to a HDF5 group. Arguments: hdf5_group: HDF5 group. optimizer: optimizer instance. """ symbolic_weights = getattr(optimizer, 'weights') if symbolic_weights: weights_group = hdf5_group.create_group('optimizer_weights') weight_names = [str(w.name).encode('utf8') for w in symbolic_weights] save_attributes_to_hdf5_group(weights_group, 'weight_names', weight_names) weight_values = K.batch_get_value(symbolic_weights) for name, val in zip(weight_names, weight_values): param_dset = weights_group.create_dataset( name, val.shape, dtype=val.dtype) if not val.shape: # scalar param_dset[()] = val else: param_dset[:] = val
def save_weights_to_hdf5_group(f, fh_dict, layers, compress=False): """Saves the weights of a list of layers to a HDF5 group. This is revised version. We split the attributes of HDF5 group into another JSON file to avoid the heading memory excessing problem. Compared to original Keras API, we need to load an extra file IO handle, fh_dict. Arguments: f: HDF5 group. fh_dict: JSON config dictionary. layers: a list of layer instances. compress: whether to compress the weights. """ compression = 'gzip' if compress else None save_attributes_to_hdf5_group(fh_dict, f.name, 'layer_names', [layer.name for layer in layers]) for layer in layers: g = f.create_group(layer.name) symbolic_weights = layer.weights weight_values = K.batch_get_value(symbolic_weights) weight_names = [] for i, (w, val) in enumerate(zip(symbolic_weights, weight_values)): if hasattr(w, 'name') and w.name: name = str(w.name) else: name = 'param_' + str(i) weight_names.append(name) save_attributes_to_hdf5_group(fh_dict, g.name, 'weight_names', weight_names) for name, val in zip(weight_names, weight_values): param_dset = g.create_dataset( name, val.shape, dtype=val.dtype, compression=(compression if val.shape else None)) if not val.shape: # scalar param_dset[()] = val else: param_dset[:] = val
def get_weights(self): params = self.weights return backend.batch_get_value(params)
def experimental_tpu_predict_loop(model, dataset, verbose=0, steps=None, callbacks=None): """Predict loop for predicting with TPU DistributionStrategy. Arguments: model: Keras Model instance. dataset: Dataset for input data. verbose: Integer, Verbosity mode 0 or 1. steps: Total number of steps (batches of samples) before declaring `_predict_loop` finished. Ignored with the default value of `None`. callbacks: List of callbacks to be called during training Returns: Array of predictions (if the model has a single output) or list of arrays of predictions (if the model has multiple outputs). """ mode = ModeKeys.PREDICT steps = training_utils.infer_steps_for_dataset(dataset, steps, steps_name='steps') dataset_fully_shaped = (distributed_training_utils. is_dataset_shape_fully_defined(dataset)) padding_handler = None if not dataset_fully_shaped: # TODO(hongjunchoi): Investigate whether operations from # PartialBatchPaddingHandler are unnecessarily pruned out # during graph optimization. padding_handler = padding_util.PartialBatchPaddingHandler( model._feed_output_shapes) batch_size, _, prefetch_buffer = input_lib._get_dataset_attributes(dataset) padding_handler.padded_batch_size = batch_size padding_handler.padding_mask = dataset.reduce(padding_handler.padding_mask, padding_handler.update_mask) dataset = dataset.map(padding_handler.pad_batch) dataset = dataset.apply(batching.unbatch()) # Upon this point, it is guaranteed that the dataset does not # have partial batches. Thus, we set `drop_remainder=True` to # get static shape information about the elements in the dataset. dataset = dataset.batch(batch_size, drop_remainder=True) if prefetch_buffer is not None: dataset = dataset.prefetch(prefetch_buffer) current_strategy = model._distribution_strategy iterator = distributed_training_utils.get_iterator(dataset, current_strategy) scope = distributed_training_utils.distributed_scope( strategy=current_strategy, learning_phase=0) scope.__enter__() out_labels = model.output_names step_fn = _make_step_fn(model, ModeKeys.PREDICT, current_strategy, out_labels) # Add initial dummy values for outputs. initial_loop_values = {} batch_dimension = distributed_training_utils.get_batch_dimension(iterator) for name, tensor in zip(model.output_names, model.outputs): # TODO(priyag): This is a workaround as we do not know the batch dimension # of the model's output at this point. shape = tensor_shape.TensorShape(tensor.shape.dims) shape.dims = [batch_dimension] + shape.dims[1:] initial_loop_values[name] = array_ops.zeros(shape, tensor.dtype) # TODO(priyag, sourabhbajaj): Support steps_per_run if/when we add outfeed. ctx = current_strategy.extended.experimental_run_steps_on_iterator( step_fn, iterator, iterations=1, initial_loop_values=initial_loop_values) predict_op = ctx.run_op output_tensors = ctx.last_step_outputs if verbose == 1: progbar = Progbar(target=steps) if model._compile_distribution: distributed_training_utils._copy_weights_to_distributed_model(model, mode) distributed_training_utils._reset_metrics(model) callbacks = cbks.configure_callbacks( callbacks, model, do_validation=False, epochs=1, steps_per_epoch=steps, verbose=verbose, count_mode='steps', mode=mode) callbacks._call_begin_hook(mode) # Since we do not know how many samples we will see, we cannot pre-allocate # the returned Numpy arrays. Instead, we store one array per batch seen # and concatenate them upon returning. unconcatenated_outs = [[] for _ in model.outputs] if steps is not None: target_steps = steps else: target_steps = np.inf current_step = 0 while current_step < target_steps: batch_logs = {'batch': current_step, 'size': 1} callbacks._call_batch_hook(mode, 'begin', current_step, batch_logs) try: _, batch_outs = K.batch_get_value([predict_op, output_tensors]) except errors.OutOfRangeError: if steps is not None: warning_msg = 'Make sure that your dataset can generate at least ' '`steps` batches (in this case, {} batches).'.format(steps) else: warning_msg = 'Number of steps ran: {} steps'.format(current_step) logging.warning('Your dataset iterator ran out of data; ' 'interrupting evaluation. ' + warning_msg) break # TODO(priyag): maybe need to unwrap the outputs first for MirroredStrategy. for i, label in enumerate(model.output_names): unconcatenated_outs[i].extend(batch_outs[label]) batch_logs = cbks.make_logs(model, batch_logs, batch_outs, mode) callbacks._call_batch_hook(mode, 'end', current_step, batch_logs) if verbose >= 1: progbar.update(current_step + 1) current_step += 1 callbacks._call_end_hook(mode) scope.__exit__(None, None, None) if len(unconcatenated_outs) == 1: prediction_result = np.concatenate(unconcatenated_outs[0], axis=0) else: prediction_result = [ np.concatenate(unconcatenated_outs[i], axis=0) for i in range(len(unconcatenated_outs)) ] if padding_handler: prediction_result = padding_handler.apply_mask(prediction_result) return prediction_result
def experimental_tpu_test_loop(model, dataset, verbose=0, steps=None, callbacks=None): """Test loop for evaluating with TPU DistributionStrategy. Arguments: model: Keras Model instance. dataset: Dataset for input data. verbose: Integer, Verbosity mode 0 or 1. steps: Total number of steps (batches of samples) before declaring predictions finished. Ignored with the default value of `None`. callbacks: List of callbacks to be called during training Returns: Scalar loss (if the model has a single output and no metrics) or list of scalars (if the model has multiple outputs and/or metrics). The attribute `model.metrics_names` will give you the display labels for the outputs. """ mode = ModeKeys.TEST current_strategy = model._distribution_strategy iterator = distributed_training_utils.get_iterator(dataset, current_strategy) steps = training_utils.infer_steps_for_dataset(dataset, steps, steps_name='steps') scope = distributed_training_utils.distributed_scope( strategy=current_strategy, learning_phase=0) scope.__enter__() out_labels = model.metrics_names def _test_step_fn(inputs): """A fn that returns output of single test step.""" inputs, targets = inputs (distribution_strategy_context.get_replica_context().merge_call( _build_model, args=(model, mode, inputs, targets))) (_, outputs, updates, _) = ( _per_device_execution_function( distributed_training_utils.get_distributed_model(model, mode), mode)) with ops.control_dependencies([updates]): return outputs test_input_data = iterator.get_next() per_replica_outputs = current_strategy.experimental_run_v2( _test_step_fn, args=(test_input_data,)) output_tensors = {} for label, output in zip(out_labels, per_replica_outputs): if label == 'loss': reduce_op = ds_reduce_util.ReduceOp.SUM else: # We reduce all other metrics using mean for now. This is temporary # workaround until new metrics are in place. reduce_op = ds_reduce_util.ReduceOp.MEAN output_tensors[label] = current_strategy.reduce(reduce_op, output) test_op = control_flow_ops.group(list(output_tensors.values())) if verbose >= 1: progbar = Progbar(target=steps) if model._compile_distribution: distributed_training_utils._copy_weights_to_distributed_model(model, mode) distributed_training_utils._reset_metrics(model) callbacks = cbks.configure_callbacks( callbacks, model, do_validation=False, epochs=1, steps_per_epoch=steps, verbose=verbose, count_mode='steps', mode=ModeKeys.TEST) callbacks._call_begin_hook(mode) outs = [0.] * len(model.metrics_names) if steps is not None: target_steps = steps else: raise ValueError('Number of steps could not be infered from the data, ' 'please pass the steps argument.') current_step = 0 while current_step < target_steps: batch_logs = {'batch': current_step, 'size': 1} callbacks._call_batch_hook(mode, 'begin', current_step, batch_logs) try: _, batch_outs = K.batch_get_value([test_op, output_tensors]) except errors.OutOfRangeError: warning_msg = 'Make sure that your dataset can generate at least ' '`steps` batches (in this case, {} batches).'.format(steps) logging.warning('Your dataset iterator ran out of data; ' 'interrupting evaluation. ' + warning_msg) target_steps = current_step break for i, label in enumerate(model.metrics_names): if i == 0: # Loss is stateless metrics. outs[i] += batch_outs[label] else: # For all stateful metrics, the aggregation is handled by mirrored vars. outs[i] = batch_outs[label] batch_logs = cbks.make_logs(model, batch_logs, outs, mode) callbacks._call_batch_hook(mode, 'end', current_step, batch_logs) if verbose == 1: progbar.update(current_step + 1) current_step += 1 if verbose >= 1: # Progress bar finishes at the end. progbar.update(target_steps) callbacks._call_end_hook(mode) scope.__exit__(None, None, None) if len(outs) >= 0: outs[0] /= (target_steps) if len(outs) == 1: return outs[0] return outs
def experimental_tpu_predict_loop(model, dataset, verbose=0, steps=None, callbacks=None): """Predict loop for predicting with TPU DistributionStrategy. Arguments: model: Keras Model instance. dataset: Dataset for input data. verbose: Integer, Verbosity mode 0 or 1. steps: Total number of steps (batches of samples) before declaring `_predict_loop` finished. Ignored with the default value of `None`. callbacks: List of callbacks to be called during training Returns: Array of predictions (if the model has a single output) or list of arrays of predictions (if the model has multiple outputs). """ mode = ModeKeys.PREDICT steps = training_utils.infer_steps_for_dataset(dataset, steps, steps_name='steps') dataset_fully_shaped = (distributed_training_utils. is_dataset_shape_fully_defined(dataset)) padding_handler = None if not dataset_fully_shaped: # TODO(hongjunchoi): Investigate whether operations from # PartialBatchPaddingHandler are unnecessarily pruned out # during graph optimization. padding_handler = padding_util.PartialBatchPaddingHandler( model._feed_output_shapes) batch_size, _, prefetch_buffer = input_lib._get_dataset_attributes(dataset) padding_handler.padded_batch_size = batch_size padding_handler.padding_mask = dataset.reduce(padding_handler.padding_mask, padding_handler.update_mask) dataset = dataset.map(padding_handler.pad_batch) dataset = dataset.apply(batching.unbatch()) # Upon this point, it is guaranteed that the dataset does not # have partial batches. Thus, we set `drop_remainder=True` to # get static shape information about the elements in the dataset. dataset = dataset.batch(batch_size, drop_remainder=True) if prefetch_buffer is not None: dataset = dataset.prefetch(prefetch_buffer) current_strategy = model._distribution_strategy iterator = distributed_training_utils.get_iterator(dataset, current_strategy) scope = distributed_training_utils.distributed_scope( strategy=current_strategy, learning_phase=0) scope.__enter__() def _predict_step_fn(inputs): """A fn that returns output of single prediction step.""" (distribution_strategy_context.get_replica_context().merge_call( _build_model, args=(model, mode, inputs))) (_, outputs, updates, _) = ( _per_device_execution_function( distributed_training_utils.get_distributed_model(model, mode), mode)) with ops.control_dependencies([updates]): return outputs # TODO(hongjunchoi): When numpy array is passed as an input to `predict()` # use numpy arrays directly to avoid cumulating unnecessary input pipeline # ops. predict_input_data = iterator.get_next() per_replica_outputs = current_strategy.experimental_run_v2( _predict_step_fn, args=(predict_input_data,)) output_tensors = distributed_training_utils.flatten_perdevice_values( current_strategy, per_replica_outputs) if verbose >= 1: progbar = Progbar(target=steps) if model._compile_distribution: distributed_training_utils._copy_weights_to_distributed_model(model, mode) distributed_training_utils._reset_metrics(model) callbacks = cbks.configure_callbacks( callbacks, model, do_validation=False, epochs=1, steps_per_epoch=steps, verbose=verbose, count_mode='steps', mode=mode) callbacks._call_begin_hook(mode) # Since we do not know how many samples we will see, we cannot pre-allocate # the returned Numpy arrays. Instead, we store one array per batch seen # and concatenate them upon returning. num_model_outputs = len(model.output_names) unconcatenated_outs = [[] for _ in range(num_model_outputs)] if steps is not None: target_steps = steps else: raise ValueError('Number of steps could not be infered from the data, ' 'please pass the steps argument.') current_step = 0 while current_step < target_steps: batch_logs = {'batch': current_step, 'size': 1} callbacks._call_batch_hook(mode, 'begin', current_step, batch_logs) try: predict_ops = control_flow_ops.group(output_tensors) _, batch_outs = K.batch_get_value([predict_ops, output_tensors]) except errors.OutOfRangeError: warning_msg = 'Make sure that your dataset can generate at least ' '`steps` batches (in this case, {} batches).'.format(steps) logging.warning('Your dataset iterator ran out of data; ' 'interrupting evaluation. ' + warning_msg) break # TODO(priyag): maybe need to unwrap the outputs first for MirroredStrategy. for i in range(num_model_outputs): output_start_index = i * current_strategy.num_replicas_in_sync output_end_index = ( output_start_index + current_strategy.num_replicas_in_sync) single_model_output = batch_outs[output_start_index:output_end_index] unconcatenated_outs[i].extend(single_model_output) batch_logs = cbks.make_logs(model, batch_logs, batch_outs, mode) callbacks._call_batch_hook(mode, 'end', current_step, batch_logs) if verbose == 1: progbar.update(current_step + 1) current_step += 1 if verbose >= 1: # Progress bar finishes at the end. progbar.update(current_step) callbacks._call_end_hook(mode) scope.__exit__(None, None, None) if len(unconcatenated_outs) == 1: prediction_result = np.concatenate(unconcatenated_outs[0], axis=0) else: prediction_result = [ np.concatenate(unconcatenated_outs[i], axis=0) for i in range(len(unconcatenated_outs)) ] if padding_handler: prediction_result = padding_handler.apply_mask(prediction_result) return prediction_result
def experimental_tpu_test_loop(model, dataset, verbose=0, steps=None, callbacks=None): """Test loop for evaluating with TPU DistributionStrategy. Arguments: model: Keras Model instance. dataset: Dataset for input data. verbose: Integer, Verbosity mode 0 or 1. steps: Total number of steps (batches of samples) before declaring predictions finished. Ignored with the default value of `None`. callbacks: List of callbacks to be called during training Returns: Scalar loss (if the model has a single output and no metrics) or list of scalars (if the model has multiple outputs and/or metrics). The attribute `model.metrics_names` will give you the display labels for the outputs. """ mode = ModeKeys.TEST current_strategy = model._distribution_strategy iterator = distributed_training_utils.get_iterator(dataset, current_strategy) steps = training_utils.infer_steps_for_dataset(dataset, steps, steps_name='steps') scope = distributed_training_utils.distributed_scope( strategy=current_strategy, learning_phase=0) scope.__enter__() out_labels = model.metrics_names step_fn = _make_step_fn(model, ModeKeys.TEST, current_strategy, out_labels) # Add initial dummy values for loss and other metric tensors. initial_loop_values = {} initial_loop_values['loss'] = constant_op.constant(1e7) for name in model.metrics_names[1:]: tensor = model._all_stateful_metrics_tensors[name] initial_loop_values[name] = array_ops.zeros(tensor.shape, tensor.dtype) # TODO(priyag): Use steps_per_run when we use new metrics as they will # allow handling metric computation at each step using variables. ctx = current_strategy.extended.experimental_run_steps_on_iterator( step_fn, iterator, iterations=1, initial_loop_values=initial_loop_values) test_op = ctx.run_op output_tensors = ctx.last_step_outputs if verbose == 1: progbar = Progbar(target=steps) if model._compile_distribution: distributed_training_utils._copy_weights_to_distributed_model(model, mode) distributed_training_utils._reset_metrics(model) callbacks = cbks.configure_callbacks( callbacks, model, do_validation=False, epochs=1, steps_per_epoch=steps, verbose=verbose, count_mode='steps', mode=ModeKeys.TEST) callbacks._call_begin_hook(mode) outs = [0.] * len(model.metrics_names) if steps is not None: target_steps = steps else: target_steps = np.inf current_step = 0 while current_step < target_steps: batch_logs = {'batch': current_step, 'size': 1} callbacks._call_batch_hook(mode, 'begin', current_step, batch_logs) try: _, batch_outs = K.batch_get_value([test_op, output_tensors]) except errors.OutOfRangeError: if steps is not None: warning_msg = 'Make sure that your dataset can generate at least ' '`steps` batches (in this case, {} batches).'.format(steps) else: warning_msg = 'Number of steps ran: {} steps'.format(current_step) logging.warning('Your dataset iterator ran out of data; ' 'interrupting evaluation. ' + warning_msg) target_steps = current_step break for i, label in enumerate(model.metrics_names): if i == 0: # Loss is stateless metrics. outs[i] += batch_outs[label] else: # For all stateful metrics, the aggregation is handled by mirrored vars. outs[i] = batch_outs[label] batch_logs = cbks.make_logs(model, batch_logs, outs, mode) callbacks._call_batch_hook(mode, 'end', current_step, batch_logs) if verbose >= 1: progbar.update(current_step + 1) current_step += 1 callbacks._call_end_hook(mode) scope.__exit__(None, None, None) if len(outs) >= 0: outs[0] /= (target_steps) if len(outs) == 1: return outs[0] return outs
def save_model(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.') from tensorflow.python.keras import __version__ as keras_version # pylint: disable=g-import-not-at-top 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: f.attrs['keras_version'] = str(keras_version).encode('utf8') f.attrs['backend'] = K.backend().encode('utf8') f.attrs['model_config'] = json.dumps( { 'class_name': model.__class__.__name__, 'config': model.get_config() }, default=serialization.get_json_type).encode('utf8') model_weights_group = f.create_group('model_weights') model_layers = model.layers save_weights_to_hdf5_group(model_weights_group, model_layers) if include_optimizer and model.optimizer: if isinstance(model.optimizer, optimizers.TFOptimizer): logging.warning( 'TensorFlow optimizers do not ' 'make it possible to access ' 'optimizer attributes or optimizer state ' 'after instantiation. ' 'As a result, we cannot save the optimizer ' 'as part of the model save file.' 'You will have to compile your model again after loading it. ' 'Prefer using a Keras optimizer instead ' '(see keras.io/optimizers).') else: f.attrs['training_config'] = json.dumps( { 'optimizer_config': { 'class_name': model.optimizer.__class__.__name__, 'config': model.optimizer.get_config() }, 'loss': model.loss, 'metrics': model.metrics, 'sample_weight_mode': model.sample_weight_mode, 'loss_weights': model.loss_weights, }, default=serialization.get_json_type).encode('utf8') # Save optimizer weights. symbolic_weights = getattr(model.optimizer, 'weights') if symbolic_weights: optimizer_weights_group = f.create_group('optimizer_weights') weight_values = K.batch_get_value(symbolic_weights) weight_names = [] for w, val in zip(symbolic_weights, weight_values): name = str(w.name) weight_names.append(name.encode('utf8')) optimizer_weights_group.attrs['weight_names'] = weight_names for name, val in zip(weight_names, weight_values): param_dset = optimizer_weights_group.create_dataset( name, val.shape, dtype=val.dtype) if not val.shape: # scalar param_dset[()] = val else: param_dset[:] = val f.flush() finally: if opened_new_file: f.close()
def experimental_tpu_fit_loop(model, dataset, epochs=100, verbose=1, callbacks=None, initial_epoch=0, steps_per_epoch=None, val_dataset=None, validation_steps=None, validation_freq=1): """Fit loop for training with TPU DistributionStrategy. Arguments: model: Keras Model instance. dataset: Dataset that returns inputs and targets epochs: Number of times to iterate over the data verbose: Integer, Verbosity mode, 0, 1 or 2 callbacks: List of callbacks to be called during training initial_epoch: Epoch at which to start training (useful for resuming a previous training run) steps_per_epoch: Total number of steps (batches of samples) before declaring one epoch finished and starting the next epoch. Ignored with the default value of `None`. val_dataset: Dataset for validation data. validation_steps: Number of steps to run validation for (only if doing validation from data tensors). Ignored with the default value of `None`. validation_freq: Only relevant if validation data is provided. Integer or `collections.Container` instance (e.g. list, tuple, etc.). If an integer, specifies how many training epochs to run before a new validation run is performed, e.g. `validation_freq=2` runs validation every 2 epochs. If a Container, specifies the epochs on which to run validation, e.g. `validation_freq=[1, 2, 10]` runs validation at the end of the 1st, 2nd, and 10th epochs. Returns: Returns `None`. Raises: ValueError: in case of invalid arguments. """ mode = ModeKeys.TRAIN # TODO(fchollet): add support for `steps_per_epoch=None` in TPU loops. current_strategy = model._distribution_strategy iterator = distributed_training_utils.get_iterator(dataset, current_strategy) steps_per_epoch = training_utils.infer_steps_for_dataset( dataset, steps_per_epoch, epochs, steps_name='steps_per_epoch') if (current_strategy.extended.steps_per_run != 1 and steps_per_epoch is None): raise ValueError('`steps_per_epoch` should be specified when calling ' '`fit` on the model with TPUStrategy when ' '`steps_per_run` != 1 .') scope = distributed_training_utils.distributed_scope( strategy=current_strategy, learning_phase=1) scope.__enter__() out_labels = model.metrics_names or [] step_fn = _make_step_fn(model, ModeKeys.TRAIN, current_strategy, out_labels) # Add initial dummy values for loss and other metric tensors. initial_loop_values = {} initial_loop_values['loss'] = constant_op.constant(1e7) for name in model.metrics_names[1:]: tensor = model._all_stateful_metrics_tensors[name] initial_loop_values[name] = array_ops.zeros(tensor.shape, tensor.dtype) use_steps = steps_per_epoch is not None if use_steps: iteration_value = min(steps_per_epoch, current_strategy.extended.steps_per_run) else: iteration_value = current_strategy.extended.steps_per_run steps_per_run = K.variable( value=iteration_value, dtype='int32', name='steps_per_run') ctx = current_strategy.extended.experimental_run_steps_on_iterator( step_fn, iterator, iterations=steps_per_run, initial_loop_values=initial_loop_values) train_op = ctx.run_op output_tensors = ctx.last_step_outputs do_validation = bool(validation_steps) if model._compile_distribution: distributed_training_utils._copy_weights_to_distributed_model(model, mode) callbacks = cbks.configure_callbacks( callbacks, model, do_validation=do_validation, epochs=epochs, steps_per_epoch=steps_per_epoch, verbose=verbose, count_mode='steps', mode=mode) # Calculate the steps each time on the device. if use_steps: steps_to_run = ([current_strategy.extended.steps_per_run] * (steps_per_epoch // current_strategy.extended.steps_per_run)) if steps_per_epoch % current_strategy.extended.steps_per_run: steps_to_run.append( steps_per_epoch % current_strategy.extended.steps_per_run) target_steps = len(steps_to_run) else: target_steps = np.inf callbacks._call_begin_hook(mode) for epoch in range(initial_epoch, epochs): distributed_training_utils._reset_metrics(model) callbacks.on_epoch_begin(epoch) epoch_logs = {} step_index = 0 prev_step_count = None current_step = 0 while current_step < target_steps: step_count = steps_to_run[current_step] if use_steps else 1 batch_logs = {'batch': step_index, 'size': 1, 'num_steps': step_count} callbacks._call_batch_hook(mode, 'begin', step_index, batch_logs) if prev_step_count is None or step_count != prev_step_count: steps_per_run.load(step_count, K.get_session()) prev_step_count = step_count try: _, outputs = K.batch_get_value([train_op, output_tensors]) except errors.OutOfRangeError: if use_steps: logging.warning('Your dataset iterator ran out of data; ' 'interrupting training. Make sure that your dataset ' 'can generate at least `steps_per_epoch * epochs` ' 'batches (in this case, %d batches).' % steps_per_epoch * epochs) else: target_steps = current_step logging.info('Dataset iterator ran out of data. Inferring the ' 'value of `steps_per_epoch` as %s .' % target_steps) distributed_training_utils.initialize_iterator(iterator, current_strategy) break batch_logs.update(outputs) callbacks._call_batch_hook(mode, 'end', step_index, batch_logs) step_index = step_index + step_count current_step += 1 if callbacks.model.stop_training: break if (do_validation and training_utils.should_run_validation(validation_freq, epoch)): logging.info('Running validation at fit epoch: %s', epoch) if model._compile_distribution: # Since we create a new clone from the original model we need to copy # the weights back to the original model before we can run validation. distributed_training_utils._copy_weights_to_original_model( model, ModeKeys.TRAIN) val_outs = experimental_tpu_test_loop( # pylint: disable=undefined-variable model, val_dataset, steps=validation_steps, verbose=verbose, callbacks=callbacks) if not isinstance(val_outs, list): val_outs = [val_outs] # Same labels assumed. for label, val_out in zip(out_labels, val_outs): epoch_logs['val_' + label] = val_out callbacks.on_epoch_end(epoch, epoch_logs) if callbacks.model.stop_training: break callbacks._call_end_hook(mode) if model._compile_distribution: # Copy the weights back from the replicated model to the original model. distributed_training_utils._copy_weights_to_original_model( model, ModeKeys.TRAIN) scope.__exit__(None, None, None) return model.history
def experimental_tpu_predict_loop(model, dataset, verbose=0, steps=None, callbacks=None): """Predict loop for predicting with TPU tf.distribute.Strategy. Arguments: model: Keras Model instance. dataset: Dataset for input data. verbose: Integer, Verbosity mode 0 or 1. steps: Total number of steps (batches of samples) before declaring `_predict_loop` finished. Ignored with the default value of `None`. callbacks: List of callbacks to be called during training Returns: Array of predictions (if the model has a single output) or list of arrays of predictions (if the model has multiple outputs). """ mode = ModeKeys.PREDICT dataset_fully_shaped = dist_utils.is_dataset_shape_fully_defined(dataset) padding_handler = None if not dataset_fully_shaped: # TODO(hongjunchoi): Investigate whether operations from # PartialBatchPaddingHandler are unnecessarily pruned out # during graph optimization. padding_handler = padding_util.PartialBatchPaddingHandler( model._feed_output_shapes) batch_size, _, prefetch_buffer = input_lib._get_dataset_attributes( dataset) padding_handler.padded_batch_size = batch_size padding_handler.padding_mask = dataset.reduce( padding_handler.padding_mask, padding_handler.update_mask) dataset = dataset.map(padding_handler.pad_batch) dataset = dataset.unbatch() # Upon this point, it is guaranteed that the dataset does not # have partial batches. Thus, we set `drop_remainder=True` to # get static shape information about the elements in the dataset. dataset = dataset.batch(batch_size, drop_remainder=True) if prefetch_buffer is not None: dataset = dataset.prefetch(prefetch_buffer) current_strategy = model._distribution_strategy iterator = dist_utils.get_iterator(dataset, current_strategy) scope = dist_utils.distributed_scope(strategy=current_strategy, learning_phase=0) scope.__enter__() def _predict_step_fn(inputs): """A fn that returns output of single prediction step.""" (distribution_strategy_context.get_replica_context().merge_call( _build_model, args=(model, mode, inputs))) (_, outputs, updates, _) = _per_replica_execution_function( dist_utils.get_distributed_model(model, mode), mode) with ops.control_dependencies([updates]): return [array_ops.identity(out) for out in outputs] # TODO(hongjunchoi): When numpy array is passed as an input to `predict()` # use numpy arrays directly to avoid cumulating unnecessary input pipeline # ops. predict_input_data = iterator.get_next() per_replica_outputs = current_strategy.run(_predict_step_fn, args=(predict_input_data, )) output_tensors = dist_utils.flatten_per_replica_values( current_strategy, per_replica_outputs) if verbose >= 1: progbar = Progbar(target=steps) if model._compile_distribution: dist_utils._copy_weights_to_distributed_model(model, mode) dist_utils._reset_metrics(model) callbacks = cbks.configure_callbacks(callbacks, model, do_validation=False, epochs=1, steps_per_epoch=steps, verbose=verbose, count_mode='steps', mode=mode) callbacks._call_begin_hook(mode) # Since we do not know how many samples we will see, we cannot pre-allocate # the returned Numpy arrays. Instead, we store one array per batch seen # and concatenate them upon returning. num_model_outputs = len(model.output_names) unconcatenated_outs = [[] for _ in range(num_model_outputs)] if steps is not None: target_steps = steps else: raise ValueError( 'Number of steps could not be inferred from the data, ' 'please pass the steps argument.') current_step = 0 while current_step < target_steps: batch_logs = {'batch': current_step, 'size': 1} callbacks._call_batch_hook(mode, 'begin', current_step, batch_logs) try: predict_ops = control_flow_ops.group(output_tensors) _, batch_outs = K.batch_get_value([predict_ops, output_tensors]) except errors.OutOfRangeError: warning_msg = ( 'Make sure that your dataset can generate at least ' '`steps` batches (in this case, {} batches).'.format(steps)) logging.warning('Your dataset iterator ran out of data; ' 'interrupting evaluation. ' + warning_msg) break # TODO(priyag): maybe need to unwrap the outputs first for MirroredStrategy. for i in range(num_model_outputs): output_start_index = i * current_strategy.num_replicas_in_sync output_end_index = (output_start_index + current_strategy.num_replicas_in_sync) single_model_output = batch_outs[ output_start_index:output_end_index] unconcatenated_outs[i].extend(single_model_output) batch_logs = cbks.make_logs(model, batch_logs, batch_outs, mode) callbacks._call_batch_hook(mode, 'end', current_step, batch_logs) if verbose == 1: progbar.update(current_step + 1) current_step += 1 if verbose >= 1: # Progress bar finishes at the end. progbar.update(current_step) callbacks._call_end_hook(mode) scope.__exit__(None, None, None) if len(unconcatenated_outs) == 1: prediction_result = np.concatenate(unconcatenated_outs[0], axis=0) else: prediction_result = [ np.concatenate(out, axis=0) for out in unconcatenated_outs ] if padding_handler: prediction_result = padding_handler.apply_mask(prediction_result) return prediction_result
def fit(self, x=None, y=None, batch_size=None, epochs=1, verbose=1, auto_switch=True, retry_fit=True, absorb=True, train_after_switch=True, callbacks=None, validation_split=0., validation_data=None, shuffle=True, class_weight=None, sample_weight=None, initial_epoch=0, steps_per_epoch=None, validation_steps=None, validation_batch_size=None, validation_freq=1, max_queue_size=10, workers=1, use_multiprocessing=False, revert_after_fit=False): """ Custom fit function for the context model auto_switch: Enable/disable autonomous context switching train_after_switch: retry_fit: Locate the next fitting context by re-performing fit. absorb: Reset the switch sequence counter upon successful training. This is mainly used to maintain switch sequencing for temporally-extended tasks revert_after_fit This is a debug parameter to revert weights after performing a fit. This is used to calculate the context deltas without incorrectly learning while auto switching is disabled """ training._keras_api_gauge.get_cell('fit').set(True) # Legacy graph support is contained in `training_v1.Model`. version_utils.disallow_legacy_graph('Model', 'fit') self._assert_compile_was_called() self._check_call_args('fit') if validation_split: # Create the validation data using the training data. Only supported for # `Tensor` and `NumPy` input. (x, y, sample_weight), validation_data = ( data_adapter.train_validation_split( (x, y, sample_weight), validation_split=validation_split, shuffle=False)) with self.distribute_strategy.scope( ), training_utils.RespectCompiledTrainableState(self): # Creates a `tf.data.Dataset` and handles batch and epoch iteration. data_handler = WindowedDataHandler( x=x, y=y, sample_weight=sample_weight, batch_size=batch_size, steps_per_epoch=steps_per_epoch, initial_epoch=initial_epoch, epochs=epochs, shuffle=shuffle, class_weight=class_weight, max_queue_size=max_queue_size, workers=workers, use_multiprocessing=use_multiprocessing, model=self) # Container that configures and calls `tf.keras.Callback`s. if not isinstance(callbacks, callbacks_module.CallbackList): callbacks = callbacks_module.CallbackList( callbacks, add_history=True, add_progbar=bool(verbose & Verbosity.Progress), model=self, verbose=verbose, epochs=epochs, steps=data_handler.inferred_steps) self.stop_training = False train_function = self.make_train_function() callbacks.on_train_begin() self.initialize_fit() # Handle fault-tolerance for multi-worker. # TODO(omalleyt): Fix the ordering issues that mean this has to # happen after `callbacks.on_train_begin`. data_handler._initial_epoch = ( self._maybe_load_initial_epoch_from_ckpt(initial_epoch)) for epoch, window_iterator in data_handler.enumerate_epochs(): self.reset_metrics() callbacks.on_epoch_begin(epoch) dataset = tf.data.Dataset.zip(next(window_iterator)) switched_during_epoch = False # Indicate if the model has attempted at least one switch during this epoch switched = True # Indicate if the model switched on the most recent fit iteration weights = backend.batch_get_value(self.trainable_variables) # Perform a 'fit call'. Assuming retry_fit, this call is re-attempted after each switch until a context fits while switched and (retry_fit or not switched_during_epoch): self.initialize_epoch(epoch) iterator = iter(dataset) # Perform a fit call with data_handler.catch_stop_iteration(): for step in data_handler.steps(): with traceme.TraceMe('TraceContext', graph_type='train', epoch_num=epoch, step_num=step, batch_size=batch_size): callbacks.on_train_batch_begin(step) tmp_logs = train_function(iterator) # Catch OutOfRangeError for Datasets of unknown size. # This blocks until the batch has finished executing. # TODO(b/150292341): Allow multiple async steps here. if not data_handler.inferred_steps: context.async_wait() logs = tmp_logs # No error, now safe to assign to logs. callbacks.on_train_batch_end(step, logs) switched = not self.update_and_switch( epoch, auto_switch, absorb, retry_fit, verbose) switched_during_epoch |= switched # If a switch occurred, we need to restore the weights if switched or (switched_during_epoch and not train_after_switch ) or revert_after_fit: backend.batch_set_value( zip(self.trainable_variables, weights)) self.reset_metrics() epoch_logs = copy.copy(logs) # Run validation. if validation_data and self._should_eval( epoch, validation_freq): val_x, val_y, val_sample_weight = ( data_adapter.unpack_x_y_sample_weight(validation_data)) val_logs = self.evaluate( x=val_x, y=val_y, sample_weight=val_sample_weight, batch_size=validation_batch_size or batch_size, steps=validation_steps, callbacks=callbacks, max_queue_size=max_queue_size, workers=workers, use_multiprocessing=use_multiprocessing, return_dict=True) val_logs = { 'val_' + name: val for name, val in val_logs.items() } epoch_logs.update(val_logs) callbacks.on_epoch_end(epoch, epoch_logs) if self.stop_training: break callbacks.on_train_end() return self.history
def save_model(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.') from tensorflow.python.keras import __version__ as keras_version # pylint: disable=g-import-not-at-top 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: f.attrs['keras_version'] = str(keras_version).encode('utf8') f.attrs['backend'] = K.backend().encode('utf8') f.attrs['model_config'] = json.dumps( { 'class_name': model.__class__.__name__, 'config': model.get_config() }, default=serialization.get_json_type).encode('utf8') model_weights_group = f.create_group('model_weights') model_layers = model.layers save_weights_to_hdf5_group(model_weights_group, model_layers) if include_optimizer and hasattr(model, 'optimizer'): if isinstance(model.optimizer, optimizers.TFOptimizer): logging.warning( 'TensorFlow optimizers do not ' 'make it possible to access ' 'optimizer attributes or optimizer state ' 'after instantiation. ' 'As a result, we cannot save the optimizer ' 'as part of the model save file.' 'You will have to compile your model again after loading it. ' 'Prefer using a Keras optimizer instead ' '(see keras.io/optimizers).') else: f.attrs['training_config'] = json.dumps( { 'optimizer_config': { 'class_name': model.optimizer.__class__.__name__, 'config': model.optimizer.get_config() }, 'loss': model.loss, 'metrics': model.metrics, 'sample_weight_mode': model.sample_weight_mode, 'loss_weights': model.loss_weights, }, default=serialization.get_json_type).encode('utf8') # Save optimizer weights. symbolic_weights = getattr(model.optimizer, 'weights') if symbolic_weights: optimizer_weights_group = f.create_group( 'optimizer_weights') weight_values = K.batch_get_value(symbolic_weights) weight_names = [] for w, val in zip(symbolic_weights, weight_values): name = str(w.name) weight_names.append(name.encode('utf8')) optimizer_weights_group.attrs[ 'weight_names'] = weight_names for name, val in zip(weight_names, weight_values): param_dset = optimizer_weights_group.create_dataset( name, val.shape, dtype=val.dtype) if not val.shape: # scalar param_dset[()] = val else: param_dset[:] = val f.flush() finally: if opened_new_file: f.close()