class_docstr: Defines how to apply the EPR algorithm.
class_docstr_tok: ["Defines","how","to","apply","the","EPR","algorithm","."]
docstr: Augments a model for training with EPR.
docstr_tok: ["Augments","a","model","for","training","with","EPR","."]
code: def get_training_model(self, model: tf.keras.Model) -> tf.keras.Model:
      """Augments a model for training with EPR."""
      if not (
          isinstance(model, tf.keras.Sequential)
          or model._is_graph_network
      ):  # pylint: disable=protected-access
          raise ValueError(
              "`model` must be either sequential or functional."
          )
  
      training_model = tf.keras.models.clone_model(
          model,
          clone_function=functools.partial(
              algorithm.create_layer_for_training, algorithm=self
          ),
      )
      training_model.build(model.input.shape)
  
      # Divide regularization weight by number of original model parameters to
      # bring it into a more standardized range.
      weight = self.regularization_weight / float(model.count_params())
  
      def regularization_loss(layer, name):
          return weight * self.regularization_loss(
              *layer.training_weights[name]
          )
  
      for layer in training_model.layers:
          if not hasattr(layer, "attr_name_map"):
              continue
          for name in layer.attr_name_map.values():
              layer.add_loss(
                  functools.partial(regularization_loss, layer, name)
              )
  
      # TODO(jballe): It would be great to be able to track the entropy losses
      # combined during training. How to do this?
      # TODO(jballe): Some models might require training log_scale weights with a
      # different optimizer/learning rate. How to do this?
      return training_model
  
imports: import functools
  from typing import Callable, List, Tuple
  import tensorflow
  import tensorflow_compression
  from tensorflow_model_optimization.python.core.common.keras.compression import algorithm
label_desc: Use image node_id 4 for calling the EPRBase obj's underlying member method code with example usage: obj.get_training_model(model) and returns: training_model, unknown