Hi Dick!



 rpfeynman:

Hi!
I’m used to the idea that backpropagation takes roughly a couple of times more than the forward prediction to run (please tell me if this is an unrealistic assumption).
Currently, I’m doing a simple linear regression where an unknown vector x is multiplied by a known  matrix A to get a prediction y'. Then I’m minimizing the MSE loss between the true y and the predicted y'.
The matrix A is quite large but sparse, and therefore is saved in a torch.sparse.FloatTensor() in the CPU.
…
The output I get is
Matrix product done in 0.36480283737182617
Back propagation done in 25.860530376434326

Is it actually possible that backpropagation takes so long to compute? Is there anything I can try to speed up, bearing in mind that A won’t fit the GPU memory?


What you report does seem surprising.  I don’t know enough
about torch.sparse and autograd to have an opinion about
whether this would be expected, but it shouldn’t have to be
this way.
You report that calculating the gradient (loss.backward())
takes about one hundred times as long as calculating the
loss itself.
Here’s how I understand what you have written:
You have a known vector of length n, y, and a sparse n x n
matrix, A.  I’ll assume that A is sparse in the conventional
sense that it has order-n non-zero elements, and that it is
stored in such a way that matrix-vector multiplication (e.g.,
A . x) can be performed in order-n time.
You also have a “hidden” length-n vector, x, such that y = A . x,
and that you are trying recover x by performing a gradient-descent
optimization.  For a given approximation to x, x’, you calculate
y’ = A . x’.  Let’s denote the “residuals” by r = y’ - y.  The MSE loss
is then sum_{i = 1, n} r_i^2.
A is fixed (and known); that is, you’re only trying to optimize for
x.  Therefore, for a given x’, you need to calculate the gradient
of loss with respect to x’, that is, grad_i = d loss / d x’_i.
Taking the derivative, you get grad_i = 2 (A-transpose . r)_i.
Because A is sparse, calculating y’, the residual r, and the
loss all take order-n time.  (If pytorch is smart enough to
remember r, you don’t have to recalculate it, but even so, it’s
only order-n.)
A is sparse, so A-transpose is too.  (Supposedly pytorch
knows this, but maybe not.)  So calculating the gradient is
also order-n, because it requires only sparse-matrix–vector
multiplication.
So I would expect loss.backward() to take about as long as
calculating loss.  (Maybe twice as long if it has to recalculate
r.  Maybe three or four times as long if it has to monkey around
traversing the autograd graph.)
You don’t say what the dimension of your matrix is, but if it
were, say, n = 100, and loss.backward() somehow wasn’t
able to use sparse-matrix multiplication, then loss.backward
would be order-n^2 (full-matrix–vector multiplication), which
would be about 100 times slower than the order-n calculation
of loss.
As written in your post, you don’t actually use y_pred in your
gradient calculation, you don’t say what the dimensionality of
your problem is, you don’t say what RF_mask is, and you don’t
say which of your tensors have requires_grad = True.
But if your actual calculation is as I laid out above, the
loss.backward() calculation could be done in order-n
and take about as long as the loss calculation.  Whether
autograd actually does it in order-n, I don’t know.
In short, if loss.backward() manages to use sparse-matrix
multiplication, it should run about as fast as calculating loss.
But if the matrix multiplication somehow gets bumped up to
full-matrix multiplication, then loss.backward() will take about
n times longer, which for “quite large” A could be a lot longer.
I, for one,  would love to hear from our pytorch / autograd
experts about whether we should expect autograd to “play
smart” with sparse matrices.
Best regards.
K. Frank