I have implemented an extended kalman filter optimizer in tensorflow. I have to calculate the jacobian instead of gradient against the prediction generated by the model. The model is generating a heatmap with detected objects from various view angles, placed on a single heatmap. I am using following loop for gradient Tape:
image, label = inputs
with tf.GradientTape(persistent=True) as tape:
predictions = self(image)
loss = self.compiled_loss(label, predictions, regularization_losses=self.losses)
Hs = tape.jacobian(predictions, self.trainable_variables)
the jacobian calculations are taking way too much time. My output i.e predictions size is 177x152 (heatmap image). I know it is creating a 26907 elements matrix. But shouldn't the tape helping the jacobian compute faster? Is my implementation for jacobian wrong? or is this simply the draw back of computing the 2nd order derivatives?