def create_layer_for_training()

in tensorflow_model_optimization/python/core/common/keras/compression/internal/optimize.py [0:0]

• 32 lines of code
• 18 McCabe index (conditional complexity)

def create_layer_for_training(layer, algorithm):
  """Internal API to create layer for training with weight compression."""

  # TODO(tfmot): move these checks to public API for
  # visibility.
  if not isinstance(algorithm, object):
    raise ValueError('`_create_layer_for_training` requires `algorithm` '
                     'to be an instantiated object, as opposed '
                     'to the class itself.')

  # Currently only supports a layer being built. The non-built
  # case may work fine as is, but it needs to be tested, together
  # with the followup API for exporting the model when the training
  # and inference graphs differ.
  if not layer.built:
    raise ValueError('`_create_layer_for_training` requires `layer` to '
                     'be built.')

  input_shape = layer.input_shape

  compressible_weights = algorithm.get_compressible_weights(layer)

  # Clone layer for two reasons:
  #
  #   1) Avoid unnecessary variable creation which undoes the benefits of
  #   compression. For instance, if we factorize `kernel` into `a` and `b`,
  #   since `a` and `b` collectively take less space than `kernel`, we
  #   no longer want to `kernel` to take up space as a variable.
  #
  #   The design depends on replacing the layer's `add_weight`
  #   method to prevent variable creation, before `add_weight` is called
  #   in the layer's `build`. Since the layer is built already, we undo
  #   this by cloning the layer.
  #
  #   2) The unoptimized layer and the optimized layer are now independent
  #   of each other and training one will not affect the other.
  #
  # TODO(tfmot): consider if it's okay to avoid this complexity during training
  # and only add it during inference, which is when model size really matters.
  # TODO(tfmot): handle custom Keras layer case.
  cloned_layer = layer.__class__.from_config(layer.get_config())

  # TODO(tfmot): We concider variable name form is layer_name/variable_name.
  layer_name = layer.name
  compressible_weights_name = []
  for compressible_weight in compressible_weights:
    name = compressible_weight.name
    if name.startswith(layer_name):
      name = name[len(layer_name)+1:]
    name = name.split(':')[0]
    compressible_weights_name.append(name)

  # TODO(tfmot): consider if this manner of handling build hinders
  # support for subclassed models in trying to set the attributes
  # that are layers while ensuring that the underlying trainable weights
  # have been created already.
  wrapped_layer = _TrainingWrapper(
      cloned_layer,
      algorithm,
      compressible_weights_name)

  if compressible_weights:
    # Set pretrained weight values.
    wrapped_layer.build(input_shape)
    training_weights = _map_to_training_weights(
        algorithm,
        layer,
        compressible_weights)
    wrapped_layer.set_weights(
        [weight.numpy() for weight in training_weights])

  return wrapped_layer