How to speed up linalg operations

john90: <p>I am training a model that needs to calculate the inverse of a [batchSize, 512, 512] matrix in the loss calculations. However, adding this step causes training to go from ~20min/epoch to ~2.5hrs/epoch. That’s even with jit scripting the function. The documentation says that almost all torch.linalg functions using CUDA will synchronize with the CPU.</p>
<p>Can anyone explain why this is necessary and, more importantly, how I can speed this up? I need to train multiple models and I don’t have 2-3 weeks for each one. My current idea is to generate an index matrix at initialization, use torch.gather to separate the minors, calculate all the determinants, and then reshape it. Is there an easier way?</p>

KFrank: <p>Hi John!</p>
<aside class="quote no-group" data-username="john90" data-post="1" data-topic="146300" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/john90/48/47499_2.png" class="avatar"> john90:</div>
<blockquote>
<p>I am training a model that needs to calculate the inverse of a [batchSize, 512, 512] matrix in the loss calculations.</p>
</blockquote>
</aside>
<p>Do I understand your use case correctly?</p>
<p>You have a batch of <code>batchSize</code> <code>512x512</code> matrices, and you need to<br>
invert each of these matrices individually.  Is this correct?</p>
<p>You say that you do this in the loss calculations.  Do you then need to<br>
<code>backward()</code> through the matrix inversions, or do the inversions not<br>
participate directly in the backpropagation?</p>
<aside class="quote no-group">
<blockquote>
<p>However, adding this step causes training to go from ~20min/epoch to ~2.5hrs/epoch.</p>
</blockquote>
</aside>
<p>This seems out of whack.  How many total samples do you have in one<br>
epoch?  That is, what is <code>batchSize</code> and how many batches do have in<br>
one epoch?</p>
<p>I can invert a 512x512 matrix in about 4 msec. on my gpu, and I can<br>
invert and backpropagate through the inversion in about 13 msec.  So<br>
you would have to have a lot of samples (about half a million) in an<br>
epoch for an inversion plus backpropagation to account for the two-hour<br>
increase in your training time.</p>
<p>Here is a timing script:</p>
<pre data-code-wrap="python"><code class="lang-python">import torch
print (torch.__version__)
print (torch.cuda.get_device_name (0))

import time

def minv_time (b, n = 1):
    print ('torch.minv():')
    print ('b.shape:', b.shape, ' b.device:', b.device)
    cnt = 0
    n_warm = 2
    for  i in range (n_warm):
        for  m in b:
            minv = torch.inverse (m)
    tm_start = time.time()
    torch.cuda.synchronize()
    for  i in range (n):
        for  m in b:
            minv = torch.inverse (m)
            cnt += 1
    torch.cuda.synchronize()
    tm = time.time() - tm_start
    print ('tm:', tm, ' cnt:', cnt, ' tm / cnt:', tm / cnt)

def minv_backward_time (b, n = 1):
    print ('torch.minv() with backward():')
    print ('b.shape:', b.shape, ' b.device:', b.device)
    b = b.clone ()
    b.requires_grad = True
    cnt = 0
    n_warm = 2
    for  i in range (n_warm):
        for  m in b:
            minv = torch.inverse (m)
            minv.sum().backward()
    tm_start = time.time()
    torch.cuda.synchronize()
    for  i in range (n):
        for  m in b:
            minv = torch.inverse (m)
            minv.sum().backward()
            cnt += 1
    torch.cuda.synchronize()
    tm = time.time() - tm_start
    print ('tm:', tm, ' cnt:', cnt, ' tm / cnt:', tm / cnt)

def mm_time (b, n = 1):
    print ('torch.mm():')
    print ('b.shape:', b.shape, ' b.device:', b.device)
    cnt = 0
    n_warm = 2
    for  i in range (n_warm):
        for  m in b:
            mm = torch.mm (m, m)
    tm_start = time.time()
    torch.cuda.synchronize()
    for  i in range (n):
        for  m in b:
            mm = torch.mm (m, m)
            cnt += 1
    torch.cuda.synchronize()
    tm = time.time() - tm_start
    print ('tm:', tm, ' cnt:', cnt, ' tm / cnt:', tm / cnt)

def bmm_time (b, n = 1):
    print ('torch.bmm():')
    print ('b.shape:', b.shape, ' b.device:', b.device)
    n_op = 0
    n_warm = 2
    for  i in range (n_warm):
        bmm = torch.bmm (b, b)
    tm_start = time.time()
    torch.cuda.synchronize()
    for  i in range (n):
        bmm = torch.bmm (b, b)
        n_op += b.size (0)
    torch.cuda.synchronize()
    tm = time.time() - tm_start
    print ('tm:', tm, ' n_op:', n_op, ' tm / n_op:', tm / n_op)

_ = torch.manual_seed (2022)

bm_cpu = torch.randn (100, 512, 512)
bm_gpu = bm_cpu.cuda()

minv_time (bm_cpu, 10)
minv_time (bm_gpu, 10)

minv_backward_time (bm_cpu, 10)
minv_backward_time (bm_gpu, 10)

mm_time (bm_cpu, 10)
mm_time (bm_gpu, 10)

bmm_time (bm_cpu, 10)
bmm_time (bm_gpu, 10)
</code></pre>
<p>And here is its output:</p>
<pre><code class="lang-plaintext">1.10.0
GeForce GTX 1050 Ti
torch.minv():
b.shape: torch.Size([100, 512, 512])  b.device: cpu
tm: 3.6806907653808594  cnt: 1000  tm / cnt: 0.0036806907653808593
torch.minv():
b.shape: torch.Size([100, 512, 512])  b.device: cuda:0
tm: 4.221258878707886  cnt: 1000  tm / cnt: 0.004221258878707885
torch.minv() with backward():
b.shape: torch.Size([100, 512, 512])  b.device: cpu
tm: 38.761736154556274  cnt: 1000  tm / cnt: 0.038761736154556276
torch.minv() with backward():
b.shape: torch.Size([100, 512, 512])  b.device: cuda:0
tm: 13.22103214263916  cnt: 1000  tm / cnt: 0.01322103214263916
torch.mm():
b.shape: torch.Size([100, 512, 512])  b.device: cpu
tm: 1.449591875076294  cnt: 1000  tm / cnt: 0.001449591875076294
torch.mm():
b.shape: torch.Size([100, 512, 512])  b.device: cuda:0
tm: 0.18179106712341309  cnt: 1000  tm / cnt: 0.00018179106712341308
torch.bmm():
b.shape: torch.Size([100, 512, 512])  b.device: cpu
tm: 1.5577073097229004  n_op: 1000  tm / n_op: 0.0015577073097229003
torch.bmm():
b.shape: torch.Size([100, 512, 512])  b.device: cuda:0
tm: 0.15090656280517578  n_op: 1000  tm / n_op: 0.00015090656280517578
</code></pre>
<p>(This script also looks at cpu and matrix-multiplication timings for context<br>
and comparison.)</p>
<aside class="quote no-group">
<blockquote>
<p>how I can speed this up?</p>
</blockquote>
</aside>
<p>First, figure out the apparent mismatch between my naive understanding<br>
of your use case and my timing results.</p>
<p>However, if you really are performing something like half a million<br>
inversions per epoch, then I doubt that you will be able to get much<br>
speed-up.</p>
<p>There could be some significant inefficiency in pytorch’s gpu handling<br>
of this sort of basic linear algebra (that you could work around), but that<br>
seems very unlikely to me.  I would trust pytorch to be highly optimized<br>
for such basic building blocks.</p>
<p>In general, linear algebra for large (and large-ish) matrices is “expensive”<br>
(matrix inversion and matrix multiplication both scale as <code>n^3</code> for <code>n x n</code><br>
matrices), but – as my timings show – for <code>512 x 512</code> matrices, it’s not<br>
that expensive.</p>
<p>Some possible suggestions:</p>
<p>Do your matrices have special structure (say, being block-diagonal) that<br>
could make them cheaper to invert?  If so, you could use a specialized<br>
(and cheaper) inversion algorithm.</p>
<p>Do you in fact need the entire inverse of the matrices?  For example, if<br>
you are solving a “single” linear algebra problem that uses the matrix<br>
(which would correspond, roughly speaking, to computing a single<br>
column of its inverse), this can be performed more cheaply than the<br>
full inversion (although it does still scale as <code>n^3</code>).</p>
<aside class="quote no-group">
<blockquote>
<p>My current idea is to generate an index matrix at initialization, use torch.gather to separate the minors, calculate all the determinants,</p>
</blockquote>
</aside>
<p>Just to be clear, if your matrix is constant across samples, or your<br>
batch of matrices is constant across batches, you should pre-compute<br>
and cache the inverse(s) once at the beginning, and then re-use the<br>
pre-computed inverse(s) over and over again for each epoch.</p>
<p>If, however, your matrices to be inverted are the result of the forward<br>
pass (so that you’re not inverting the same matrix multiple times), then<br>
trying to implement your own inversion routine by calculating the<br>
determinants of the minors (the <em>cofactors</em>) won’t do you any good,<br>
and will almost certainly slow things down (significantly).</p>
<p>Good luck!</p>
<p>K. Freak</p>

john90: <p>Hi KFrank,</p>
<p>Thank you for your reply! You hit the nail on the head. I have half a million samples in each epoch and I need to backpropagate through the inversions. However, I don’t need the full inverse, but it’s complicated. I’m implementing the following equation, where <strong>B</strong> is a 60x512 matrix, <strong>D</strong> is 512x512, and lambda is a predicted value. Even though the indices will always be the same, the matrices come from the forward pass. But, since the bulk of the second matrix is 0-padded, I can cut down on how much of the inverse I need to calculate. So I need to include all the columns but only the first 60 rows.</p>
<p><img src="https://discuss.pytorch.org/uploads/default/original/3X/a/5/a535bda37029b967561881ec0ad16be33778e134.png" alt="image" data-base62-sha1="nzvV91OzCrSxNYq3Tt9w70Y3XEM" width="285" height="163"></p>
<p>Unfortunately, I haven’t found an easy way to implement this. I have code to generate the indices and calculate everything, but a matrix like that is hard to handle given the memory constraints.</p>
<p>Maybe my best bet is a feature request that allows us to specify partial inverses. Or maybe they can find an alternative to synchronizing with the CPU.</p>
<p>Thanks again for your detailed response!</p>
<p>Best,<br>
John</p>

KFrank: <p>Hi John!</p>
<p>There are a number of things that could be said about this, but I<br>
don’t really understand what you are doing.</p>
<aside class="quote no-group" data-username="john90" data-post="3" data-topic="146300" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/john90/48/47499_2.png" class="avatar"> john90:</div>
<blockquote>
<p><strong>B</strong> is a 60x512 matrix, <strong>D</strong> is 512x512, and lambda is a predicted value.</p>
</blockquote>
</aside>
<p>Let me call the matrix you are trying to “invert” <code>M</code>, and write it out<br>
as <code>[[B], [sqrt (lambda) * D]]</code>.</p>
<p>I therefore suppose that <code>M</code> is a <code>572 x 512</code> matrix.  (If not, what is it?)<br>
As such, it is not square, and does not have an inverse.</p>
<p>Are you imagining that the “inverse” is some sort of <em>pseudoinverse?</em><br>
(This could make sense, depending on what you are doing.)   I could<br>
speculate, but better that you clarify the details of what your equation<br>
means.</p>
<aside class="quote no-group">
<blockquote>
<p>Maybe my best bet is a feature request that allows us to specify partial inverses.</p>
</blockquote>
</aside>
<p>I don’t know whether it would be relevant to your use case, but you can<br>
think of the <a href="https://pytorch.org/docs/1.11/generated/torch.lu.html" rel="noopener nofollow ugc">LU decomposition</a> of a matrix as being a “partial inverse” in<br>
that you can use it to solve multiple equations, <code>M @ x = v</code>, with the same<br>
matrix, <code>M</code>, but different vector right-hand-sides, <code>v</code>, without  computing the<br>
full inverse of <code>M</code>, but also not starting from scratch each time.</p>
<aside class="quote no-group">
<blockquote>
<p>Or maybe they can find an alternative to synchronizing with the CPU.</p>
</blockquote>
</aside>
<p>I don’t understand why you are so fixated on this cpu synchronization.</p>
<p>In my timing tests (that seem to correspond to your two-hour per epoch<br>
run-times) my gpu was nearly saturated at 98%.  If you need two hours<br>
of gpu time, that’s how much time you need (even if the synchronization<br>
did represent an unnecessary cpu inefficiency, which I doubt, because I<br>
trust pytorch on these basic building blocks).</p>
<p>Best.</p>
<p>K. Frank</p>