Python tensorflow.function方法代码示例

# 需要导入模块: import tensorflow [as 别名]
# 或者: from tensorflow import function [as 别名]
def compute_model(self, inputs, variables, training):
    """Compute the model for a set of inputs and variables.

    Parameters
    ----------
    inputs: list of tensors
      the inputs to the model
    variables: list of tensors
      the values to use for the model's variables.  This might be the actual
      variables (as returned by the MetaLearner's variables property), or
      alternatively it might be the values of those variables after one or more
      steps of gradient descent for the current task.
    training: bool
      indicates whether the model is being invoked for training or prediction

    Returns
    -------
    (loss, outputs) where loss is the value of the model's loss function, and
    outputs is a list of the model's outputs
    """
    raise NotImplemented("Subclasses must implement this") 

# 需要导入模块: import tensorflow [as 别名]
# 或者: from tensorflow import function [as 别名]
def _create_gradient_fn(self, variables):
    """Create a function that computes gradients and applies them to the model.
    Because of the way TensorFlow function tracing works, we need to create a
    separate function for each new set of variables.
    """

    @tf.function(experimental_relax_shapes=True)
    def apply_gradient_for_batch(inputs, labels, weights, loss):
      with tf.GradientTape() as tape:
        outputs = self.model(inputs, training=True)
        if isinstance(outputs, tf.Tensor):
          outputs = [outputs]
        if self._loss_outputs is not None:
          outputs = [outputs[i] for i in self._loss_outputs]
        batch_loss = loss(outputs, labels, weights)
      if variables is None:
        vars = self.model.trainable_variables
      else:
        vars = variables
      grads = tape.gradient(batch_loss, vars)
      self._tf_optimizer.apply_gradients(zip(grads, vars))
      self._global_step.assign_add(1)
      return batch_loss

    return apply_gradient_for_batch 

# 需要导入模块: import tensorflow [as 别名]
# 或者: from tensorflow import function [as 别名]
def set_dataset_ready(self, dataset, set_map: bool = True, set_batch: bool = True):
		dataset = dataset.prefetch(tf.data.experimental.AUTOTUNE)

		if set_map:
			try:
				if self.mode == "triplet":
					dataset = dataset.map(self.mapper_triplet, tf.data.experimental.AUTOTUNE)
				elif self.mode == "softmax" or self.mode == "sparse softmax":
					dataset = dataset.map(self.mapper_softmax, tf.data.experimental.AUTOTUNE)
				elif self.mode == "arcface":
					dataset = dataset.map(self.mapper_arcface, tf.data.experimental.AUTOTUNE)
				else:
					raise Exception(f"There is no mapping function for {self.mode}, please fix.")
			except ValueError:
				raise Exception(f"You must set dataset for {self.mode} if you want to use.")

		if set_batch:
			dataset = dataset.batch(self.batch_size)

		print("Dataset ready for stream.")

		return dataset 

# 需要导入模块: import tensorflow [as 别名]
# 或者: from tensorflow import function [as 别名]
def test_none_state_equal_to_initial_state(self):
    """Tests that not providing state is the same as initial_state."""
    x = tf.constant(1.0)
    encoder = simple_encoder.SimpleEncoder(
        core_encoder.EncoderComposer(
            test_utils.PlusOneOverNEncodingStage()).make(),
        tf.TensorSpec.from_tensor(x))

    state = encoder.initial_state()
    stateful_iteration = _make_iteration_function(encoder)

    @tf.function
    def stateless_iteration(x):
      encoded_x, _ = encoder.encode(x)
      decoded_x = encoder.decode(encoded_x)
      return encoded_x, decoded_x

    _, encoded_x_stateful, decoded_x_stateful, _ = self.evaluate(
        stateful_iteration(x, state))
    encoded_x_stateless, decoded_x_stateless = self.evaluate(
        stateless_iteration(x))

    self.assertAllClose(encoded_x_stateful, encoded_x_stateless)
    self.assertAllClose(decoded_x_stateful, decoded_x_stateless) 

# 需要导入模块: import tensorflow [as 别名]
# 或者: from tensorflow import function [as 别名]
def compute_loss(self, pi_, z_, policy_head, value_head):

        # loss
        with tf.name_scope("loss"):
            policy_loss = tf.keras.losses.categorical_crossentropy(y_true=pi_, y_pred=policy_head, from_logits=True)
            policy_loss = tf.reduce_mean(policy_loss)

            value_loss = tf.keras.losses.mean_squared_error(z_, value_head)
            value_loss = tf.reduce_mean(value_loss)
            # summary_ops_v2.scalar('mse_loss', value_loss)

            regularizer = tf.keras.regularizers.l2(self.c_l2)
            regular_variables = self.model.trainable_variables
            l2_loss = self.apply_regularization(regularizer, regular_variables)

            #             self.loss = value_loss - policy_loss + l2_loss
            self.loss = value_loss + policy_loss + l2_loss
            # summary_ops_v2.scalar('loss', self.loss)

        return self.loss

    # TODO(yashkatariya): Add tf.function when b/123315763 is resolved
    # @tf.function 

# 需要导入模块: import tensorflow [as 别名]
# 或者: from tensorflow import function [as 别名]
def test_optimizer(should_quantize):
    x = QuantizedVariable.from_variable(get_var(1.0), quantizer=lambda x: -x)
    opt = tf.keras.optimizers.SGD(1.0)

    def loss():
        with context.quantized_scope(should_quantize):
            return x + 1.0

    @tf.function
    def f():
        opt.minimize(loss, var_list=[x])

    f()
    if should_quantize:
        assert evaluate(x) == 2.0
        with context.quantized_scope(should_quantize):
            assert evaluate(x) == -2.0
    else:
        assert evaluate(x) == 0.0 

# 需要导入模块: import tensorflow [as 别名]
# 或者: from tensorflow import function [as 别名]
def mlp(input_tensor, layers, activ_fn=tf.nn.relu, layer_norm=False):
    """
    Create a multi-layer fully connected neural network.

    :param input_tensor: (tf.placeholder)
    :param layers: ([int]) Network architecture
    :param activ_fn: (tf.function) Activation function
    :param layer_norm: (bool) Whether to apply layer normalization or not
    :return: (tf.Tensor)
    """
    output = input_tensor
    for i, layer_size in enumerate(layers):
        output = tf.layers.dense(output, layer_size, name='fc' + str(i))
        if layer_norm:
            output = tf.contrib.layers.layer_norm(output, center=True, scale=True)
        output = activ_fn(output)
    return output 

# 需要导入模块: import tensorflow [as 别名]
# 或者: from tensorflow import function [as 别名]
def evaluate(self, dataset, epoch):
        """ evaluate the model """
        loss_metric = tf.keras.metrics.Mean(name="AverageLoss")
        loss, metrics = None, None
        evaluate_step = self.evaluate_step
        if self.hparams.enable_tf_function:
            logging.info("please be patient, enable tf.function, it takes time ...")
            evaluate_step = tf.function(evaluate_step, input_signature=self.sample_signature)
        self.model.reset_metrics()  # init metric.result() with 0
        for batch, samples in enumerate(dataset):
            samples = self.model.prepare_samples(samples)
            loss, metrics = evaluate_step(samples)
            if batch % self.hparams.log_interval == 0:
                logging.info(self.metric_checker(loss, metrics, -2))
            total_loss = sum(list(loss.values())) if isinstance(loss, dict) else loss
            loss_metric.update_state(total_loss)
        logging.info(self.metric_checker(loss_metric.result(), metrics, evaluate_epoch=epoch))
        self.model.reset_metrics()
        return loss_metric.result(), metrics 

# 需要导入模块: import tensorflow [as 别名]
# 或者: from tensorflow import function [as 别名]
def train(self, dataset, total_batches=-1):
        """ Update the model in 1 epoch """
        train_step = self.train_step
        if self.hparams.enable_tf_function:
            logging.info("please be patient, enable tf.function, it takes time ...")
            train_step = tf.function(train_step, input_signature=self.sample_signature)
        for batch, samples in enumerate(dataset.take(total_batches)):
            # train 1 step
            samples = self.model.prepare_samples(samples)
            loss, metrics = train_step(samples)
            # Horovod: broadcast initial variable states from rank 0 to all other processes.
            # This is necessary to ensure consistent initialization of all workers when
            # training is started with random weights or restored from a checkpoint.
            #
            # Note: broadcast should be done after the first gradient step to ensure optimizer
            # initialization.
            if batch == 0:
                hvd.broadcast_variables(self.model.trainable_variables, root_rank=0)
                hvd.broadcast_variables(self.optimizer.variables(), root_rank=0)
            if batch % self.hparams.log_interval == 0 and hvd.rank() == 0:
                logging.info(self.metric_checker(loss, metrics))
                self.model.reset_metrics() 

# 需要导入模块: import tensorflow [as 别名]
# 或者: from tensorflow import function [as 别名]
def parallel_projection2d(volume, geometry):
    """
    Wrapper function for making the layer call.
    Args:
        volume:     Input volume to project.
        geometry:   Corresponding GeometryParallel2D Object defining parameters.
    Returns:
            Initialized lme_custom_ops.parallel_projection2d layer.
    """
    batch = np.shape(volume)[0]
    return pyronn_layers.parallel_projection2d(volume,
                                               projection_shape=geometry.sinogram_shape,
                                               volume_origin=np.broadcast_to(geometry.volume_origin, [batch, *np.shape(geometry.volume_origin)]),
                                               detector_origin=np.broadcast_to(geometry.detector_origin, [batch, *np.shape(geometry.detector_origin)]),
                                               volume_spacing=np.broadcast_to(geometry.volume_spacing, [batch, *np.shape(geometry.volume_spacing)]),
                                               detector_spacing=np.broadcast_to(geometry.detector_spacing, [batch, *np.shape(geometry.detector_spacing)]),
                                               ray_vectors=np.broadcast_to(geometry.ray_vectors, [batch, *np.shape(geometry.ray_vectors)])) 

# 需要导入模块: import tensorflow [as 别名]
# 或者: from tensorflow import function [as 别名]
def fan_projection2d(volume, geometry):
    """
    Wrapper function for making the layer call.
    Args:
        volume:     Input volume to project.
        geometry:   Corresponding GeometryFan2D Object defining parameters.
    Returns:
            Initialized lme_custom_ops.fan_projection2d layer.
    """
    batch = np.shape(volume)[0]
    return pyronn_layers.fan_projection2d(volume,
                                          projection_shape=geometry.sinogram_shape,
                                          volume_origin=np.broadcast_to(geometry.volume_origin, [batch, *np.shape(geometry.volume_origin)]),
                                          detector_origin=np.broadcast_to(geometry.detector_origin, [batch, *np.shape(geometry.detector_origin)]),
                                          volume_spacing=np.broadcast_to(geometry.volume_spacing, [batch, *np.shape(geometry.volume_spacing)]),
                                          detector_spacing=np.broadcast_to(geometry.detector_spacing, [batch, *np.shape(geometry.detector_spacing)]),
                                          source_2_isocenter_distance=np.broadcast_to(geometry.source_isocenter_distance, [batch, *np.shape(geometry.source_isocenter_distance)]),
                                          source_2_detector_distance=np.broadcast_to(geometry.source_detector_distance, [batch, *np.shape(geometry.source_detector_distance)]),
                                          central_ray_vectors=np.broadcast_to(geometry.central_ray_vectors, [batch, *np.shape(geometry.central_ray_vectors)])) 

# 需要导入模块: import tensorflow [as 别名]
# 或者: from tensorflow import function [as 别名]
def build_single_label_dataset(dataset, label_key, desired_label):
  """Build a new dataset that only yields examples with a particular label.

  This can be used for creating pathological non-iid (in label space) datasets.

  Args:
    dataset: the base `tf.data.Dataset` that yields examples that are structures
      of string key -> tensor value pairs.
    label_key: the `str` key that holds the label for the example.
    desired_label: the label value to restrict the resulting dataset to.

  Returns:
    A `tf.data.Dataset` that is composed of only examples that have a label
    matching `desired_label`.
  """

  @tf.function
  def _select_on_label(example):
    return example[label_key] == desired_label

  return dataset.filter(_select_on_label) 

# 需要导入模块: import tensorflow [as 别名]
# 或者: from tensorflow import function [as 别名]
def from_tff_result(cls, anon_tuple):
    # TODO(b/123092620): These conversions should not be needed.
    return assert_no_anon_tuples(
        cls(
            generator_weights=list(anon_tuple.generator_weights),
            discriminator_weights=list(anon_tuple.discriminator_weights),
            counters=anon_tuple.counters._asdict(),
            # TODO(b/123092620): Using _asdict(recursive=True) is a work-around
            # which at least gets rid of AnonymousTuples to allow the use of
            # tf.nest. However, really these should be the appropriate
            # namedtuple types expected by the TF Privacy code. This
            # means that in some cases ServerState.dp_averaging_state
            # needs dict-style access, and sometimes attribute-style.
            # However, since this is really opaque state, this only comes up
            # in the test.
            dp_averaging_state=anon_tuple.dp_averaging_state._asdict(
                recursive=True)))


# Set cmp=False to get a default hash function for tf.function. 

# 需要导入模块: import tensorflow [as 别名]
# 或者: from tensorflow import function [as 别名]
def _apply_gradients(self, inputs, actions_matrix, discounted_rewards,
                       advantages, action_prob):
    """Compute the gradient of the loss function for a batch and update the model."""
    vars = self._model.model.trainable_variables
    with tf.GradientTape() as tape:
      outputs = self._model.model(inputs)
      loss = self._model._loss_fn(outputs, [actions_matrix],
                                  [discounted_rewards, advantages, action_prob])
    gradients = tape.gradient(loss, vars)
    self._model._tf_optimizer.apply_gradients(zip(gradients, vars)) 

# 需要导入模块: import tensorflow [as 别名]
# 或者: from tensorflow import function [as 别名]
def predict(self, state, use_saved_states=True, save_states=True):
    """Compute the policy's output predictions for a state.

    If the policy involves recurrent layers, this method can preserve their internal
    states between calls.  Use the use_saved_states and save_states arguments to specify
    how it should behave.

    Parameters
    ----------
    state: array or list of arrays
      the state of the environment for which to generate predictions
    use_saved_states: bool
      if True, the states most recently saved by a previous call to predict() or select_action()
      will be used as the initial states.  If False, the internal states of all recurrent layers
      will be set to the initial values defined by the policy before computing the predictions.
    save_states: bool
      if True, the internal states of all recurrent layers at the end of the calculation
      will be saved, and any previously saved states will be discarded.  If False, the
      states at the end of the calculation will be discarded, and any previously saved
      states will be kept.

    Returns
    -------
    the array of action probabilities, and the estimated value function
    """
    results = self._predict_outputs(state, use_saved_states, save_states)
    return [results[i] for i in (self._action_prob_index, self._value_index)]