Function which outputs tensor by reference

username: qualix
cooked: <p>I’m struggling to write a function ConvToToeplitzTr() which takes a conv2d’s weight matrix W as the input, constructs a new matrix using its entries (a Toeplitz matrix representation of the given conv2d, transposed) <em>by reference</em> and returns this new matrix M. Both W and M are then used to calculate the output of the layer.</p>
<p>My goal is for autograd to actually treat the variables in the original conv2d’s weight matrix W, as well as the variables in the output matrix of ConvToToeplitzTr() M as the same and correspondingly accumulate the gradients for each of these variables from both W and M.</p>
<p>The following minimal example already fails to return a tensor by reference:</p>
<pre><code>import torch

def foo(x):
    return torch.stack([x,x])

y = torch.tensor(4)
z = foo(y)
print(z) # [4, 4]
y += 1
print(z) # [4, 4]. Want: [5,5]</code></pre>

username: albanD
cooked: <p>Hi,</p>
<p>I think the notion of “reference” that you want here is what we call view: another Tensor that looks at the same memory.<br>
If you want to do that, you will need to make that you only use functions that return views. You can find more details and the list of these functions here:  <a href="https://pytorch.org/docs/stable/tensor_view.html" class="inline-onebox">Tensor Views — PyTorch 1.7.0 documentation</a></p>

username: qualix
cooked: <p>Hello albanD,</p>
<p>thank you very much for your reply.</p>
<p>I get the idea, but my specific case seems to be somewhat more complicated.</p>
<p>Here’s my actual code, note esp. the assignment to <code>m</code> inside the loops before the <code>return</code> statement:</p>
<pre><code>def convmatrix2d(kernel, image_shape):
    # kernel: (out_channels, in_channels, kernel_height, kernel_width, ...)
    # image: (in_channels, image_height, image_width, ...)
    assert image_shape[0] == kernel.shape[1]
    assert len(image_shape[1:]) == len(kernel.shape[2:])
    result_dims = torch.tensor(image_shape[1:]) - torch.tensor(kernel.shape[2:]) + 1
    m = torch.zeros((
        kernel.shape[0],
        *result_dims,
        *image_shape
    ))
    for i in range(m.shape[1]):
        for j in range(m.shape[2]):
            m[:,i,j,:,i:i+kernel.shape[2],j:j+kernel.shape[3]] = kernel.view(kernel.shape) # Previously: ... = kernel
    return m.flatten(0, len(kernel.shape[2:])).flatten(1)

# In forward() of my network:
self.convtest = nn.Conv2d(1, 1, kernel_size=3, padding=1)
output_toeplitz = convmatrix2d(self.convtest.weight, [1, 5, 5])
</code></pre>
<p>Even with the addition of <code>.view()</code>, <code>output_toeplitz</code> variable does <em>not</em> change if we now to manually change some values in <code>self.convtest.weight</code>.</p>

username: albanD
cooked: <p>If you want <code>output_toeplitz</code> to change when you update <code>self.convtest.weight</code>, then you have to make sure that you only use ops that are views to get from to the other.<br>
Here you create a Tensor full of zeros that has brand new memory. So this one can’t be a view of <code>kernel</code>.</p>
<p>You might want to use the <code>Unfold</code> class we provide to do that to get a view properly: <a href="https://pytorch.org/docs/stable/generated/torch.nn.Unfold.html?highlight=im2col" class="inline-onebox">Unfold — PyTorch 1.7.0 documentation</a></p>

username: qualix
cooked: <p>My overarching goal, born from a mathematical algorithm, is to calculate a matrix-vector product of the transposed <a href="https://stackoverflow.com/questions/16798888/2-d-convolution-as-a-matrix-matrix-multiplication/44039201#44039201" rel="noopener nofollow ugc">Toeplitz matrix</a> (associated to a convolution kernel) with a vector. It is important that gradients w.r.t. the kernel can be calculated afterwards.</p>
<p>The way I’m trying to go (maybe there’s a better way in PyTorch?) is: somehow <em>build a Toeplitz matrix of a conv2d kernel in such a manner that it still works with autograd</em>, then transpose it, then <em>perform matrix-vector multiplication</em> with it. With the code I’ve posted above I’m able to set up the Toeplitz matrix, however, it indeed duplicates the tensors, which breaks autograd.</p>
<p>I’ve taken a look at Fold/Unfold, thank you. However, I’m not sure how they can help me, since I’m trying to <em>manipulate a convolution kernel</em>. I’m not trying to manipulate an input image (in a block-based manner, which Unfold/Fold facilitate). Maybe I’m not seeing how to use Unfold to achieve my goal.</p>
<p>What would be a good way to do that in PyTorch?</p>

username: albanD
cooked: <aside class="quote no-group" data-username="qualix" data-post="5" data-topic="108933">
<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/qualix/48/33459_2.png" class="avatar"> qualix:</div>
<blockquote>
<p>it indeed duplicates the tensors, which breaks autograd.</p>
</blockquote>
</aside>
<p>I am not sure this is true.<br>
You don’t use detach or no grad so the gradients will be properly tracked through your <code>convmatrix2d</code> function (even though the output matrix does not share its memory with the input, if you backprop, the gradients will flow back all the way to the input.</p>

username: qualix
cooked: <p>Looks like you’re right! For <code>output_toeplitz()</code> defined as above:</p>
<pre><code>convtest = nn.Conv2d(1, 1, kernel_size=3, padding=1)
output_toeplitz = convmatrix2d(convtest.weight, [1, 5, 5])

for _ in range(3):
    t = torch.matmul(output_toeplitz,torch.randn(25,1))
    print("grad after matmul with output_toeplitz:")
    print(convtest.weight.grad) # None in 1st loop, zeros(3,3) afterwards
    t.sum().backward(retain_graph = True)
    print("grad after .backward():")
    print(convtest.weight.grad) # Nonzero updated gradient
    convtest.zero_grad()
</code></pre>
<p>A quick side question: why doesn’t this code work without <code>retain_graph = True</code> in <code>backward()</code>? I get <code>RuntimeError: Trying to backward through the graph a second time, but the saved intermediate results have already been freed.</code></p>

username: albanD
cooked: <p>This happens because you run the <code>convmatrix2d</code> only once for the 3 iterations of your loop. Which means that this part of the autograd graph is shared by all 3 iterations of the loop.<br>
But if you don’t state retain_graph=True, the graph is freed during backward. And here, since you have a shared part in the graph, you don’t want it to be freed.</p>
<p>Also if you update the value of <code>convtest.weight</code> (via an optimizer for example), you want to recall <code>convmatrix2d</code> to make sure you reconstruct the matrix based on the new content of <code>convtest.weight</code>!</p>

username: qualix
cooked: <p>Thank you, very insightful!</p>