Python v2.GradientTape方法代码示例

# 需要导入模块: from tensorflow.compat import v2 [as 别名]
# 或者: from tensorflow.compat.v2 import GradientTape [as 别名]
def testGrad(self):

    def f(a, b):
      return tf_np.sum(tf_np.sqrt(tf_np.exp(a)) + b)

    g = extensions.grad(f)

    def compare(a, b):
      with tf.GradientTape() as tape:
        tape.watch(a.data)
        r = f(a, b)
      expected = tape.gradient(r.data, a.data)
      self.assertAllEqual(expected, g(a, b))

    shape = [10]
    a = tf_np.random.randn(*shape)
    b = tf_np.random.randn(*shape)
    compare(a, b) 

# 需要导入模块: from tensorflow.compat import v2 [as 别名]
# 或者: from tensorflow.compat.v2 import GradientTape [as 别名]
def hessian(function: Callable[[Parameters], tf.Tensor],
            parameters: Parameters) -> Parameters:
  """Computes the Hessian of a given function.

  Useful for testing, although scales very poorly.

  Args:
    function: A function for which we want to compute the Hessian.
    parameters: Parameters with respect to the Hessian should be computed.

  Returns:
    A tensor or list of tensors of same nested structure as `Parameters`,
      representing the Hessian.
  """
  with tf.GradientTape() as outer_tape:
    with tf.GradientTape() as inner_tape:
      value = function(parameters)
    grads = inner_tape.gradient(value, parameters)
    grads = tensor_list_util.tensor_list_to_vector(grads)
  return outer_tape.jacobian(grads, parameters) 

# 需要导入模块: from tensorflow.compat import v2 [as 别名]
# 或者: from tensorflow.compat.v2 import GradientTape [as 别名]
def test_multiple_state_vars(self):
    x = tf.constant([3.0, 4.0])
    y = tf.constant([5.0, 6.0])
    z = tf.constant([7.0, 8.0])
    alpha = tf.constant(2.0)
    beta = tf.constant(1.0)

    with tf.GradientTape(persistent=True) as tape:
      tape.watch([alpha, beta])
      def body(i, state):
        x, y, z = state
        k = tf.cast(i + 1, tf.float32)
        return [x * alpha - beta, y * k * alpha * beta, z * beta + x]
      out = for_loop(body, [x, y, z], [alpha, beta], 3)

    with self.subTest("independent_vars"):
      grad = tape.gradient(out[1], alpha)
      self.assertAllEqual(792, grad)
    with self.subTest("dependent_vars"):
      grad = tape.gradient(out[2], beta)
      self.assertAllEqual(63, grad) 

# 需要导入模块: from tensorflow.compat import v2 [as 别名]
# 或者: from tensorflow.compat.v2 import GradientTape [as 别名]
def test_batching(self):
    x = tf.constant([[3.0, 4.0], [30.0, 40.0]])
    y = tf.constant([[5.0, 6.0], [50.0, 60.0]])
    z = tf.constant([[7.0, 8.0], [70.0, 80.0]])
    alpha = tf.constant(2.0)
    beta = tf.constant(1.0)

    with tf.GradientTape(persistent=True) as tape:
      tape.watch([alpha, beta])
      def body(i, state):
        x, y, z = state
        k = tf.cast(i + 1, tf.float32)
        return [x * alpha - beta, y * k * alpha * beta, z * beta + x]
      out = for_loop(body, [x, y, z], [alpha, beta], 3)
    with self.subTest("independent_vars"):
      grad = tape.gradient(out[1], alpha)
      self.assertAllEqual(8712, grad)
    with self.subTest("dependent_vars"):
      grad = tape.gradient(out[2], beta)
      self.assertAllEqual(783, grad) 

# 需要导入模块: from tensorflow.compat import v2 [as 别名]
# 或者: from tensorflow.compat.v2 import GradientTape [as 别名]
def test_with_xla(self):
    @tf.function
    def fn():
      x = tf.constant([[3.0, 4.0], [30.0, 40.0]])
      y = tf.constant([[7.0, 8.0], [70.0, 80.0]])
      alpha = tf.constant(2.0)
      beta = tf.constant(1.0)
      with tf.GradientTape(persistent=True) as tape:
        tape.watch([alpha, beta])
        def body(i, state):
          del i
          x, y = state
          return [x * alpha - beta, y * beta + x]
        out = for_loop(body, [x, y], [alpha, beta], 3)
      return tape.gradient(out[1], beta)

    grad = self.evaluate(tf.xla.experimental.compile(fn))[0]
    self.assertAllEqual(783, grad) 

# 需要导入模块: from tensorflow.compat import v2 [as 别名]
# 或者: from tensorflow.compat.v2 import GradientTape [as 别名]
def train_step(model, loss_fn, optimizer_fn, metric, image, label):
  """Perform one training step for the model.

  Args:
    model: Keras model to train.
    loss_fn: Loss function to use.
    optimizer_fn: Optimizer function to use.
    metric: keras.metric to use.
    image: Tensor of training images of shape [batch_size, 28, 28, 1].
    label: Tensor of class labels of shape [batch_size].
  """
  with tf.GradientTape() as tape:
    preds = model(image)
    label_onehot = tf.one_hot(label, 10)
    loss_ = loss_fn(label_onehot, preds)
  grads = tape.gradient(loss_, model.trainable_variables)
  optimizer_fn.apply_gradients(zip(grads, model.trainable_variables))
  metric(loss_) 

# 需要导入模块: from tensorflow.compat import v2 [as 别名]
# 或者: from tensorflow.compat.v2 import GradientTape [as 别名]
def run_one_epoch(
        self, minibatches: Iterable[np.ndarray], training: bool = False,
    ):
        total_loss, num_samples, num_tokens, num_correct_tokens = 0.0, 0, 0, 0
        for step, minibatch_data in enumerate(minibatches):
            with tf.GradientTape() as tape:
                model_outputs = self.compute_logits(minibatch_data, training=training)
                result = self.compute_loss_and_acc(model_outputs, minibatch_data)

            total_loss += result.token_ce_loss
            num_samples += minibatch_data.shape[0]
            num_tokens += result.num_predictions
            num_correct_tokens += result.num_correct_token_predictions

            if training:
                gradients = tape.gradient(
                    result.token_ce_loss, self.trainable_variables
                )
                self.optimizer.apply_gradients(zip(gradients, self.trainable_variables))

            print(
                "   Batch %4i: Epoch avg. loss: %.5f || Batch loss: %.5f | acc: %.5f"
                % (
                    step,
                    total_loss / num_samples,
                    result.token_ce_loss,
                    result.num_correct_token_predictions
                    / (float(result.num_predictions) + 1e-7),
                ),
                end="\r",
            )
        print("\r\x1b[K", end="")
        return (
            total_loss / num_samples,
            num_correct_tokens / (float(num_tokens) + 1e-7),
        ) 

# 需要导入模块: from tensorflow.compat import v2 [as 别名]
# 或者: from tensorflow.compat.v2 import GradientTape [as 别名]
def train(self, x, y, learning_rate=0.01):
    """Runs a single training pass.

    Args:
      x: 2-d array of size batch_size x image_size.
      y: 2-d array of size batch_size x num_classes in one-hot notation.
      learning_rate: The learning rate.
    """
    x = np.array(x, copy=False)
    y = np.array(y, copy=False)

    def mean_squared_error(x, y):
      diff = x - y
      return np.sum(diff * diff) / len(x)

    wb_tensors = [p.data for p in self.weights + self.biases]
    with tf.GradientTape() as g:
      g.watch(wb_tensors)
      loss = mean_squared_error(self.forward(x), y)
    gradients = g.gradient(loss.data, wb_tensors)
    gradients = [np.asarray(grad) for grad in gradients]

    new_weights_and_biases = []
    for v, dv in zip(self.weights + self.biases, gradients):
      new_weights_and_biases.append(v - learning_rate * dv)

    total_len = len(new_weights_and_biases)
    self.weights = new_weights_and_biases[:total_len // 2]
    self.biases = new_weights_and_biases[total_len // 2:] 

# 需要导入模块: from tensorflow.compat import v2 [as 别名]
# 或者: from tensorflow.compat.v2 import GradientTape [as 别名]
def testVjp(self, has_aux):
    x_shape = (tf.TensorShape([10]), tf.TensorShape([1, 10]))
    y_shape = (tf.TensorShape([]))
    dtype = np.float32

    def f(a, b):
      y = tf_np.sum(tf_np.sqrt(tf_np.exp(a)) + b)
      if has_aux:
        return y, tf_np.asarray(1)
      else:
        return y

    rng = tf.random.Generator.from_seed(1234)
    x, dy_list = tf.nest.map_structure(lambda shape: uniform(rng, shape, dtype),
                                       [x_shape, [y_shape] * 2])
    tf_x = to_tf(x)
    outputs = extensions.vjp(f, *x, has_aux=has_aux)
    if has_aux:
      y, vjp, aux = outputs
    else:
      y, vjp = outputs
    with tf.GradientTape(persistent=True) as tape:
      tape.watch(tf_x)
      outputs = f(*x)
      if has_aux:
        expected_y, expected_aux = outputs
        self.assertAllClose(to_tf(expected_aux), to_tf(aux))
      else:
        expected_y = outputs
    self.assertAllClose(to_tf(expected_y), to_tf(y))
    for dy in dy_list:
      expected_dx = tape.gradient(
          to_tf(expected_y), tf_x, output_gradients=to_tf(dy))
      self.assertAllClose(expected_dx, to_tf(vjp(dy))) 

# 需要导入模块: from tensorflow.compat import v2 [as 别名]
# 或者: from tensorflow.compat.v2 import GradientTape [as 别名]
def test_setitem(self):
    # Single integer index.
    a = array_ops.array([1., 2., 3.])
    b = array_ops.array(5.)
    c = array_ops.array(10.)

    tensors = [arr.data for arr in [a, b, c]]
    with tf.GradientTape() as g:
      g.watch(tensors)
      a[1] = b + c
      loss = array_ops.sum(a)

    gradients = g.gradient(loss.data, tensors)
    self.assertSequenceEqual(
        array_ops.array(gradients[0]).tolist(), [1., 0., 1.])
    self.assertEqual(array_ops.array(gradients[1]).tolist(), 1.)
    self.assertEqual(array_ops.array(gradients[2]).tolist(), 1.)

    # Tuple index.
    a = array_ops.array([[[1., 2.], [3., 4.]], [[5., 6.],
                                                [7., 8.]]])  # 2x2x2 array.
    b = array_ops.array([10., 11.])

    tensors = [arr.data for arr in [a, b]]
    with tf.GradientTape() as g:
      g.watch(tensors)
      a[(1, 0)] = b
      loss = array_ops.sum(a)

    gradients = g.gradient(loss.data, tensors)
    self.assertSequenceEqual(
        array_ops.array(gradients[0]).tolist(),
        [[[1., 1.], [1., 1.]], [[0., 0.], [1., 1.]]])
    self.assertEqual(array_ops.array(gradients[1]).tolist(), [1., 1.]) 

# 需要导入模块: from tensorflow.compat import v2 [as 别名]
# 或者: from tensorflow.compat.v2 import GradientTape [as 别名]
def _hessian_vector_product(
    function: Callable[[Parameters], tf.Tensor],
    parameters: Parameters,
    v: Parameters) -> Parameters:
  """Computes Hessian-vector products.

  Computes the product H.v where v is an arbitrary vector and H is the Hessian
  of a function evaluated at `parameters`.

  The result is the same as if the Hessian was computed explicitly and
  multiplied the vector. However, this function uses the autograd in backward
  then forward mode in order to compute this Hessian vector product without
  having to explicitly compute the Hessian.

  Args:
    function: A (twice) differentiable function that takes as input a tensor or
      a list of tensors and returns a scalar.
    parameters: The parameters with respect to which we want to compute the
      Hessian for the hessian vector product.
    v: An arbitrary vector or list of vectors of the same nested structure as
      `parameters`.

  Returns:
    A vector or list of vectors of the same nested structure as
      `parameters`, equal to H.v.
  """
  with tf.autodiff.ForwardAccumulator(
      primals=parameters, tangents=v) as acc:
    with tf.GradientTape() as tape:
      tape.watch(parameters)
      value = function(parameters)
    backward = tape.gradient(value, parameters)
  return acc.jvp(backward) 

# 需要导入模块: from tensorflow.compat import v2 [as 别名]
# 或者: from tensorflow.compat.v2 import GradientTape [as 别名]
def test_simple_grad_wrt_parameter(self):
    x = tf.constant([3.0])
    sigma = tf.constant(2.0)

    with tf.GradientTape() as tape:
      tape.watch(sigma)
      def body(i, state):
        del i
        x = state[0]
        return [x * sigma]
      out = for_loop(body, [x], [sigma], 3)[0]

    grad = tape.gradient(out, sigma)
    self.assertAllEqual(36, grad) 

# 需要导入模块: from tensorflow.compat import v2 [as 别名]
# 或者: from tensorflow.compat.v2 import GradientTape [as 别名]
def test_simple_grad_wrt_initial_state(self):
    x = tf.constant([3.0])
    sigma = tf.constant(2.0)

    with tf.GradientTape() as tape:
      tape.watch(x)
      def body(i, state):
        del i
        x = state[0]
        return [x * sigma]
      out = for_loop(body, [x], [sigma], 3)[0]

    grad = tape.gradient(out, x)
    self.assertAllEqual([8], grad) 

# 需要导入模块: from tensorflow.compat import v2 [as 别名]
# 或者: from tensorflow.compat.v2 import GradientTape [as 别名]
def test_shapes(self, state_dims, num_params, times):
    # Checks that the loop can handle various shapes and outputs correct shapes.
    def test_with_batch_shape(batch_shape):
      initial_state = [tf.ones(shape=batch_shape + (d,)) for d in state_dims]
      params = [tf.constant(1.0) for _ in range(num_params)]
      with tf.GradientTape(persistent=True) as tape:
        tape.watch(initial_state)
        tape.watch(params)
        def body(i, state):
          del i
          if not params:
            return state
          sum_params = tf.add_n(params)
          state = [s * sum_params for s in state]
          return state
        final_state = for_loop(body, initial_state, params, times)

      for s_in in initial_state:
        for s_out in final_state:
          grad = tape.gradient(s_out, s_in)
          self.assertAllEqual(s_in.shape, grad.shape)

      for p in params:
        for s_out in final_state:
          grad = tape.gradient(s_out, p)
          self.assertAllEqual([], grad.shape)

    with self.subTest("no_batch"):
      test_with_batch_shape(batch_shape=())
    with self.subTest("simple_batch"):
      test_with_batch_shape(batch_shape=(5,))
    with self.subTest("complex_batch"):
      test_with_batch_shape(batch_shape=(2, 8, 3)) 

# 需要导入模块: from tensorflow.compat import v2 [as 别名]
# 或者: from tensorflow.compat.v2 import GradientTape [as 别名]
def _jacobian_wrt_parameter(y, param, tape):
  """Computes a Jacobian w.r.t. a parameter."""
  # For input shapes (b, dy), yields shape (b, dy, 1) (1 is added for
  # convenience elsewhere).
  # To avoid having to broadcast param to y's shape, we need to take a forward
  # gradient.
  with tf.GradientTape() as w_tape:
    w = tf.zeros_like(y)
    w_tape.watch(w)
    vjp = tape.gradient(y, param, output_gradients=w)
  if vjp is None:  # Unconnected.
    return tf.expand_dims(tf.zeros_like(y), axis=-1)
  return tf.expand_dims(w_tape.gradient(vjp, w), axis=-1)