Memory issues with iterative algorithm

username: mithrandir
cooked: <p>So I am training a network for image reconstruction as part of an unrolled optimization scheme, where first I apply a neural network to an image to get a denoised image, then I apply a data consistency step where I do a conjugate gradient algorithm with respect to the denoised image. I apply . So the forward function of the method in pseudo code is</p>
<p>x=init_guess</p>
<p>for i in num_iterations:<br>
denoised_x=neural_network(x)</p>
<pre><code> data_consistent_x=conj_grad(data,denoised_x)

x=data_consistent_x
</code></pre>
<p>However, while training the inclusion of the conj grad step is adding hugely to the GPU memory used (without: 2GB, with: 45GB).</p>
<p>This conjugate gradient algorithm (<a href="https://en.wikipedia.org/wiki/Conjugate_gradient_method" class="inline-onebox" rel="noopener nofollow ugc">Conjugate gradient method - Wikipedia</a>) is an iterative method involving matrix operations. so there is another for loop inside.</p>
<p>Below is the function I am using, where init_image is x, and the masked_kspace,sens_mps,mask are the data.</p>
<p>def conjg_grad_correction(self,init_image,masked_kspace,sens_mps,mask):</p>
<pre><code>    r=self.my_rhs(masked_kspace,sens_mps,mask,init_image)

    p=r

    r_complex=torch.complex(r[...,0],r[...,1])
    p_complex=torch.complex(p[...,0],p[...,1])
    
    rtr=(torch.conj(r_complex)*r_complex).sum()
    x_complex=torch.zeros_like(r_complex)

    for i in range(10):
        Ap=self.myAtAp(p,sens_mps,mask)
        Ap_complex=torch.complex(Ap[...,0],Ap[...,1])
        #p_complex=torch.complex(p[...,0],p[...,1])
        alpha=rtr/((torch.conj(p_complex)*Ap_complex).sum())
        #print(alpha)
        x_complex=x_complex+alpha*p_complex
        r_complex=r_complex-alpha*Ap_complex
        rTrnew=(torch.conj(r_complex)*r_complex).sum()
        #print(rTrnew)
        beta=rTrnew/rtr
        #print(beta)
        p_complex=r_complex+beta*p_complex
        p=torch.view_as_real(p_complex)

        rtr=rTrnew
    
    return torch.view_as_real(x_complex)
</code></pre>
<p>I am rather new to pytorch, so I am wondering if it is possible to reduce the memory used during training by modifying this conj grad implementation?</p>
<p>Thanks in advance!</p>

username: KFrank
cooked: <p>Hi Mithrandir!</p>
<p>I don’t understand your use case or follow what you are doing in your<br>
code, so I can’t give you a concrete answer.  But I do have some<br>
comments, below.</p>
<aside class="quote no-group quote-modified" data-username="mithrandir" data-post="1" data-topic="111663" 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/mithrandir/48/34593_2.png" class="avatar"> mithrandir:</div>
<blockquote>
<p>then I apply a data consistency step where I do a conjugate gradient algorithm</p>
<p>…</p>
<p>However, while training the inclusion of the conj grad step is adding hugely to the GPU memory used (without: 2GB, with: 45GB).</p>
<p>This conjugate gradient algorithm […] is an iterative method involving matrix operations.</p>
</blockquote>
</aside>
<p>First, it is perfectly fine to use a loop or iterative algorithm in your<br>
forward pass.  As long as the calculations inside of your loop use<br>
(differentiable) pytorch tensor operations (that support autograd),<br>
autograd and backpropagation will work properly.</p>
<p>As you’ve seen, however, each iteration through the loop will add<br>
another “layer” to the “computation graph,” consuming memory, and<br>
taking time when you backpropagate, perhaps to the extent that you<br>
run out of memory or the computation becomes impractically slow.</p>
<p>Depending on your use case, you may be able to turn off autograd for<br>
your iterative forward function, and write your own corresponding<br>
backward function that requires less memory (and/or time).</p>
<p>The “baseline” use of the conjugate-gradient algorithm is as a linear<br>
solver where the linear operation <code>A . x</code> (with <code>A</code> a matrix, and <code>x</code> a<br>
vector) is given implicitly somehow, or where <code>A</code> is sparse so that<br>
computing <code>A . x</code> is cheaper than computing a full matrix-vector<br>
product or it is impractical to store the full <code>A^-1</code> in memory.</p>
<p>Let’s say that this is your use case, so that you are solving <code>A . x = b</code><br>
for <code>x</code>, and you require gradients of the computed <code>x</code> with respect to the<br>
elements of the vector <code>b</code>.</p>
<p>Then <code>x = A^-1 . b</code>, so that the Jacobean (of <code>x</code> with respect to <code>b</code>) is<br>
<code>A^-1</code>.  In backpropagation you typically compute the dot product of<br>
such a Jacobean with a vector of numerical gradients that are being<br>
backpropagated from later stages of your network.  Let’s call this<br>
vector <code>g</code> and the result of the dot product <code>r = A^-1 . g</code>.</p>
<p><code>r</code> can now be obtained by solving <code>A . r = g</code> for <code>r</code> for which you can<br>
again employ the conjugate-gradient algorithm, taking advantage of<br>
any sparsity in <code>A</code> (or if you need to because the linear operation<br>
<code>A . x</code> is only given implicitly).</p>
<p>You can also write your own backward function, potentially using the<br>
conjugate-gradient algorithm, in the case that you require the gradients<br>
of the computed <code>x</code> with respect to the elements of the matrix <code>A</code>.  Doing<br>
so is relatively straightforward, but may or may not be helpfully efficient<br>
or practical depending on the details of the sparsity structure of <code>A</code>.</p>
<p>Best.</p>
<p>K. Frank</p>