TensorFlow.js中的tf.grad()函数

在TensorFlow.js的核心API中，tf.grad()函数是非常重要的一个组成部分。

该函数可以用于计算函数关于变量的梯度，并返回一个新的张量作为结果。

语法

tf.grad(f, x, dy?) => (x: tf.Tensor[]) => tf.Tensor[]

参数

• f: (x: tf.Tensor[]) => tf.Tensor：输入函数，接收张量数组作为输入，并返回一个张量作为结果。
• x: tf.Tensor[]：输入张量数组。
• dy?: tf.Tensor[]：可选参数，形状与f(x)相同的张量数组。如果提供dy，则计算dy与结果的乘积dt/dy，否则将dt/dy设置为1。

返回值

• (x: tf.Tensor[]) => tf.Tensor[]：返回一个函数，该函数可以计算f(x)关于x的梯度。

示例

以下是一个简单示例，展示了如何使用tf.grad()函数计算函数y = 3x² + 2x + 1关于x的导数：

// 定义函数
const f = (x) => tf.tidy(() => {
  const y = x.square().mul(3).add(x.mul(2)).add(1);
  return y;
});

// 定义输入张量x
const x = tf.tensor1d([1, 2, 3, 4]);

// 计算y关于x的导数
const grad = tf.grad(f);
const dydx = grad(x);

// 输出结果
dydx.print();

输出结果：

Tensor
    [ 8, 16, 24, 32 ]
    shape: [ 4 ]
    dtype: float32

备注

• tf.grad()函数返回的是一个计算梯度的函数，需要将输入张量传递给该函数才能计算梯度。
• tf.grad()函数只能用于计算具有可求导性质的函数关于变量的梯度。
• 在使用tf.grad()函数时，应该将计算图包含在tidy()函数内，以便在计算完梯度后自动释放不需要的中间张量。