Tensorflow 加速雅可比求值

我想计算和评估一个NN的雅可比矩阵关于输入。我并不太关心构造雅可比矩阵所需的时间，我更关心的是雅可比矩阵的计算

weights = {
    'w1': tf.Variable(tf.random_normal([num_input, num_hidden_1])),
    'w2': tf.Variable(tf.random_normal([num_hidden_1, num_hidden_2])),
    'w_final': tf.Variable(tf.random_normal([num_hidden_2, 1]))
}
biases = {
    'b1': tf.Variable(tf.random_normal([num_hidden_1])),
    'b2': tf.Variable(tf.random_normal([num_hidden_2])),
    'b_final': tf.Variable(tf.random_normal([num_hidden_2])),
}

def g(x):
    layer_1 = tf.nn.sigmoid(tf.add(tf.matmul(x, weights['w1']),
                                   biases['b1']))
    layer_2 = tf.nn.sigmoid(tf.add(tf.matmul(layer_1, weights['w2']),
                                   biases['b2']))

    final = tf.add(tf.matmul(layer_2, weights['w_final']),
                                   biases['b_final'])
    return final

现在是雅可比矩阵的计算

# https://github.com/tensorflow/tensorflow/issues/675
def jacobian(y_flat, x):
    n = y_flat.shape[0]

    loop_vars = [
        tf.constant(0, tf.int32),
        tf.TensorArray(tf.float32, size=n),
    ]

    _, jacobian = tf.while_loop(
        lambda j, _: j < n,
        lambda j, result: (j+1, result.write(j, tf.gradients(y_flat[j], x))),
        loop_vars)

    return jacobian.stack()

这是我的输出：

0 jacobian constructed
6 Seconds: (500, 1, 500, 784)
4 Seconds: (500, 1, 500, 784)
5 Seconds: (500, 1, 500, 784)
5 Seconds: (500, 1, 500, 784)
6 Seconds: (500, 1, 500, 784)
3 Seconds: (500, 1, 500, 784)
4 Seconds: (500, 1, 500, 784)
3 Seconds: (500, 1, 500, 784)
3 Seconds: (500, 1, 500, 784)
3 Seconds: (500, 1, 500, 784)

---

有没有办法加快速度？我看不出为什么这个速度这么慢，梯度下降在理论上不是很相似吗？

嗯，最简单的解决方案毕竟是最快的

def jacobian(y, x):
    with tf.name_scope("jacob"):
        grads = tf.stack([tf.gradients(yi, x)[0] for yi in tf.unstack(y, axis=1)],
                        axis=2)
        return grads

def jacobian(y, x):
    with tf.name_scope("jacob"):
        grads = tf.stack([tf.gradients(yi, x)[0] for yi in tf.unstack(y, axis=1)],
                        axis=2)
        return grads