Tensorflow 使用 GradientTape 时矩阵大小不兼容

Question

I am trying to run code that previously worked on tensorflow 2.2.0 on version 2.4.0-rc0 for apple silicon (using python 3.8), but it is now generating the following error regarding the matrix dimensions:

tensorflow.python.framework.errors_impl.InvalidArgumentError: GetOutputShape: Matrix size-incompatible: In[0]: [256,4], In[1]: [4,400]

I am using nested gradient tapes to compute the gradient of my MLP model wrt the inputs (which form part of the loss), after which I compute the gradient of the loss wrt the trainable variables as below:

    def get_grad_and_loss(self, x, y):
        with tf.GradientTape(persistent=True) as gl_tape:
            gl_tape.watch(x)

            with tf.GradientTape(persistent=True) as l_tape:
                l_tape.watch(x)
                y_pred = self.call(x)

            grad_mat = l_tape.gradient(y_pred, x)
            loss = tf.reduce_mean(tf.math.square(y_pred - y[:, tf.newaxis])) + tf.reduce_mean(tf.maximum(0, -1 * (grad_mat[:, 0])))

        g = gl_tape.gradient(loss, self.trainable_weights)

        return g, loss

In words I am computing the MSE and trying to force the sign of the gradient to be positive (as a soft constraint). I have read through the documentation on gradient tape and as I understand it, setting persistent=True should allow me to recompute gradients freely. As a side note my code works fine if I omit the nested gradient tape and simply use the MSE metric, so I don't think the issue lies anywhere else in the code. Any pointers would be much appreciated, thanks in advance:)

Answer 1

You seem to have confusion over which gradient tape watches which variables. I suggest to make sure that the tapes watch over different vars. Presently they both watch x . Most likely you need to enter gl_tape.watch(self.trainable_weights) . There are examples out there with 2 gradient tapes working together. Check them out.