How to preserve autograd of tensor after .detach() and processing it?

Chuong_Vo: <p>Hello!</p>
<p>In the work that I’m doing, after the first <code>conv2d()</code> layer, the output is converted to numpy array to do some processing using <code>.detach()</code>. During this process, the new output will be 3 times bigger and then it is converted back to the tensor to be used as a input for the next <code>conv2d()</code> layer.</p>
<p>Is there anyway of getting the gradient back to the new tensor?</p>
<p>Note: The new tensor’s values are just copied from the old tensor.</p>
<p>Thank you guys very much!!!</p>

JuanFMontesinos: <p>Hi,</p>
<aside class="quote no-group" data-username="Chuong_Vo" data-post="1" data-topic="68483">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/c/b2d939/48.png" class="avatar"> Chuong_Vo:</div>
<blockquote>
<p>Note: The new tensor’s values are just copied from the old tensor.</p>
</blockquote>
</aside>
<p>Could you deeply explain what are you doing?<br>
What I’ve understood is that you have a tensor T, you convert it to numpy (T_numpy) but then the network keeps using T.<br>
If that is the case you don’t need to struggle, you can keep using T as if nothing happened.</p>
<pre><code class="lang-auto">import torch

a=torch.rand(10).requires_grad_()

b=a.sqrt().mean()

c=b.detach()
b.backward()
print(b.grad_fn)
print(c.grad_fn)
</code></pre>
<pre><code class="lang-auto">&lt;MeanBackward0 object at 0x7fba8eefdcc0&gt;
None
</code></pre>
<p>In case you want to modify T according to what you have done in numpy, the easiest way is to reimplement that in pytorch.<br>
Else, you can create a nn.Module which implements a proper backward.<br>
The last thing you can do is to in-place modify values of T.</p>
<pre><code class="lang-auto">import torch

a=torch.rand(10).requires_grad_()

b=a.sqrt().mean()
print(b)
b.data=torch.tensor(5).float()
print(b)
c=b.detach()
b.backward()
print(b.grad_fn)
print(c.grad_fn)
</code></pre>
<p>Which would go against DL theory</p>

Chuong_Vo: <p>Hi, thank you for replying!<br>
In my work , I want to change the output of the <code>conv2d</code> layer before putting it through another <code>conv2d</code> layer. But I had to use a nested for loop for this process this output and it takes me too much time. So I have to convert it into numpy and turn the loop into <code>numba</code> to reduce the processing time. Then I can convert the T_numpy back to tensor and use this new tensor to continue in the next <code>conv2d</code> layer.<br>
I saw in <a href="https://discuss.pytorch.org/t/should-it-really-be-necessary-to-do-var-detach-cpu-numpy/35489/3">here</a> that Mr <a class="mention" href="/u/alband">@albanD</a> stated that:</p>
<blockquote>
<p>People not very familiar with  <code>requires_grad</code>  and cpu/gpu Tensors might go back and forth with numpy. For example doing pytorch → numpy → pytorch and backward on the last Tensor. This will backward without issue but not all the way to the first part of the code and won’t raise any error.</p>
</blockquote>
<p>So is it possible to just leave it that way and continue training? Is it possible when I do that many times?</p>

JuanFMontesinos: <p>I would suggest to ask for help to optimize the nested loop rather than doing it. Realize you can write your custom c++ module (properly done) if you can’t really optimize that nested loop. You can “run” that but consider pytorch won’t be aware of all those operations, thus, will pass wrong gradients…</p>
<p>Lastly just some random links</p><aside class="quote quote-modified" data-post="1" data-topic="13923">
  <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/malmaud/48/3420_2.png" class="avatar">
    <a href="https://discuss.pytorch.org/t/accelerating-layers-with-numba/13923">Accelerating layers with Numba</a> 
  </div>
  <blockquote>
    Hi all, 
Has anyone had success integrating functions decorated with numba.cuda.jit into a PyTorch workflow? <a href="http://pytorch.org/about/" rel="noopener nofollow ugc">http://pytorch.org/about/</a> mentions using your ‘favorite libraries…such as Numba’ but I’m not clear on how to actually do that. 
Of course it’s no problem on the CPU to decorate a layer with numba.jit and convert inputs and outputs to and from numpy arrays, although even that doesn’t seem particularly elegant because of all the  numpy/Tensor conversions required (I’m aware they are zero-ov…
  </blockquote>
</aside>
<p>
And I know there is a library which creates cuda kernels to perform optimized operations for a given code, speeding it up, but i forgot the name <img src="https://discuss.pytorch.org/images/emoji/apple/confused.png?v=12" title=":confused:" class="emoji" alt=":confused:" loading="lazy" width="20" height="20"></p>

Chuong_Vo: <p>Thank you very much for the information!<br>
I think I will open a new topic about optimizing the nested loop in pytorch version.</p>
<p>If you can recall the name of the mentioned library please inform me! I’m desperately short in time for my college graduation project. Any help would definitely mean a huge favor to me! Thank you!!!</p>

albanD: <p>I think you misunderstood what I meant there: the code will run without crashing. But the gradients <strong>won’t</strong> flow back all the way to the original pytorch code !<br>
Using <code>.detach()</code> will prevent the gradients from flowing back.</p>

Chuong_Vo: <p>Hi <a class="mention" href="/u/alband">@albanD</a>,<br>
So if I just run it then the training will not be accurate right?<br>
If so, is there anyway that I can attach the gradient of the original tensor to the new tensor converted from numpy array?<br>
Thank you!</p>

albanD: <blockquote>
<p>So if I just run it then the training will not be accurate right?</p>
</blockquote>
<p>Well the gradients won’t be correct.</p>
<blockquote>
<p>If so, is there anyway that I can attach the gradient of the original tensor to the new tensor converted from numpy array?</p>
</blockquote>
<p>No. For the gradients to be properly computed you need to only use pytorch’s functions.<br>
Note that you can add new elementary function where you define <strong>both</strong> the forward and backward if you don’t know how to write the forward using pytorch’s functions (see tutorial <a href="https://pytorch.org/docs/stable/notes/extending.html">here</a>). But that will be more complex as you need to implement the backward yourself.</p>

Chuong_Vo: <p>Hi <a class="mention" href="/u/alband">@albanD</a>,<br>
I’m trying to keep the gradient of the tensor in the way below, but the outcome is not as expected. Please kindly check it out if this has any problem that can be solved. Thank you!!!</p>
<pre><code class="lang-auto">class S2ConvNet_original(nn.Module):

    def __init__(self):
        super(S2ConvNet_original, self).__init__()

        f1 = 21
        f2 = 40

        self.conv1 = nn.Conv2d(1, f1, 3, stride=3, padding=0)
        self.upsample1 = nn.UpsamplingNearest2d(scale_factor=3)
        self.conv2 = nn.Conv2d(f1, f2, 3, stride=3, padding=0)
        self.upsample2 = nn.UpsamplingNearest2d(scale_factor=3)
        self.fc1 = nn.Linear(in_features=40*60*60, out_features=10)


    def forward(self, x, coor):
        x1 = self.upsample1(x)
        x = preprocess60x60(x, x1, coor)
        x = self.conv1(x)
        x = F.relu(x)
        x2 = self.upsample2(x)
        x = preprocess60x60_21chan(x, x2, coor)
        x = self.conv2(x)
        x = F.relu(x)
        x = self.fc1(x.reshape(x.shape[0], -1))

        return x
"""
This part is for the processing functions
"""
@jit(nopython=False)
def change_patches(patches, image, coor):
    
    a = 0
    for i in range(60): #2160
        for j in range(60): #3840
            coordinate = [i,j]
            indices = np.matrix(coor[(coordinate[0],coordinate[1])])
            patches[a,0,2,1] = image[0,0,indices[1,0],indices[0,0]]
            patches[a,0,2,2] = image[0,0,indices[1,1],indices[0,1]]
            patches[a,0,1,2] = image[0,0,indices[1,2],indices[0,2]]
            patches[a,0,0,2] = image[0,0,indices[1,3],indices[0,3]]
            patches[a,0,0,1] = image[0,0,indices[1,4],indices[0,4]]
            patches[a,0,0,0] = image[0,0,indices[1,5],indices[0,5]]
            patches[a,0,1,0] = image[0,0,indices[1,6],indices[0,6]]
            patches[a,0,2,0] = image[0,0,indices[1,7],indices[0,7]]
            a += 1
    return(patches)

def preprocess60x60(images,newbatch, coor):
    DEVICE = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    
    for batch in range(images.shape[0]):
        image = torch.zeros(1,1,60,60)
        image[0,:,:,:] = images[batch,:,:,:]
        
        kc, kh, kw = 1, 3, 3  # kernel size
        dc, dh, dw = 1, 1, 1  # stride
    
        pad = (1,1,1,1)
        paddedimage = F.pad(image, pad).to(DEVICE)        
    
        patches = paddedimage.unfold(1, kc, dc).unfold(2, kh, dh).unfold(3, kw, dw).to(DEVICE)
        unfold_shape = patches.size()
        
        patches = patches.contiguous().view(-1, kc, kh, kw).to(DEVICE)
        numpypatches = patches.cpu().numpy()
        numpyimage = image.cpu().numpy()
        patches = change_patches(numpypatches, numpyimage, coor)
    
        # Reshape back
        patches_orig = torch.from_numpy(patches).view(unfold_shape)
        output_c = unfold_shape[1] * unfold_shape[4]
        output_h = unfold_shape[2] * unfold_shape[5]
        output_w = unfold_shape[3] * unfold_shape[6]
        patches_orig = patches_orig.permute(0, 1, 4, 2, 5, 3, 6).contiguous()
        patches_orig = patches_orig.view(1, output_c, output_h, output_w)
        newbatch[batch,:,:,:] = patches_orig
    return newbatch

def preprocess60x60_21chan(images, newbatch, coor):
    DEVICE = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

    newimage = torch.zeros(1,images.shape[1],180,180).to(DEVICE)

    for batch in range(images.shape[0]):
        image = torch.zeros(1,images.shape[1],60,60)
        image[0,:,:,:] = images[batch,:,:,:]
        for channel in range(21):
            
            imagei = torch.zeros(1,1,60,60)
            imagei[0,0,:,:] = image[0,channel,:,:]

            kc, kh, kw = 1, 3, 3  # kernel size
            dc, dh, dw = 1, 1, 1  # stride
        
            pad = (1,1,1,1)
            paddedimage = F.pad(imagei, pad).to(DEVICE)        
        
            patches = paddedimage.unfold(1, kc, dc).unfold(2, kh, dh).unfold(3, kw, dw).to(DEVICE)
            unfold_shape = patches.size()
            
            patches = patches.contiguous().view(-1, kc, kh, kw)
            
            numpypatches = patches.detach().cpu().numpy()
            numpyimage = imagei.detach().cpu().numpy()
            patches = change_patches(numpypatches, numpyimage, coor)
    
            # Reshape back
            patches_orig = torch.from_numpy(patches).view(unfold_shape)
            output_c = unfold_shape[1] * unfold_shape[4]
            output_h = unfold_shape[2] * unfold_shape[5]
            output_w = unfold_shape[3] * unfold_shape[6]
            patches_orig = patches_orig.permute(0, 1, 4, 2, 5, 3, 6).contiguous()
            patches_orig = patches_orig.view(1, output_c, output_h, output_w)
            newimage[0,channel,:,:] = patches_orig
        
        newbatch[batch,:,:,:] = newimage
    return newbatch

"""
The main part for common training and testing
"""
def main(network):
    train_loader, test_loader, train_dataset, _ = load_data( MNIST_PATH, BATCH_SIZE)
    with open('dict60x60.pkl', 'rb') as fp:
        coor = pickle.load(fp)
    if network == 'original':
        classifier = S2ConvNet_original()
    elif network == 'deep':
        classifier = S2ConvNet_deep()
    else:
        raise ValueError('Unknown network architecture')
    classifier.to(DEVICE)

    print("#params", sum(x.numel() for x in classifier.parameters()))

    criterion = nn.CrossEntropyLoss()
    criterion = criterion.to(DEVICE)

    optimizer = torch.optim.Adam(
        classifier.parameters(),
        lr=LEARNING_RATE)


    for epoch in range(NUM_EPOCHS):
        for i, (images, labels) in enumerate(train_loader):
            classifier.train()

            images = images.to(DEVICE)
            #print(images.size())
            
            labels = labels.to(DEVICE)
        
            optimizer.zero_grad()
            outputs = classifier(images, coor)
            loss = criterion(outputs, labels)
            loss.backward()

            optimizer.step()

            print('\rEpoch [{0}/{1}], Iter [{2}/{3}] Loss: {4:.4f}'.format(
                epoch+1, NUM_EPOCHS, i+1, len(train_dataset)//BATCH_SIZE,
                loss.item()), end="")
        print("")
        correct = 0
        total = 0
        for images, labels in test_loader:

            classifier.eval()

            with torch.no_grad():
                images = images.to(DEVICE)
                labels = labels.to(DEVICE)

                outputs = classifier(images, coor)
                _, predicted = torch.max(outputs, 1)
                total += labels.size(0)
                correct += (predicted == labels).long().sum().item()

        print('Test Accuracy: {0}'.format(100 * correct / total))

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument("--network",
                        help="network architecture to use",
                        default='original',
                        choices=['original', 'deep'])
    args = parser.parse_args()

    main(args.network)
</code></pre>

albanD: <p>You do convert a lot of thing to numpy inside your forward function (in preprocessXXX in particular). So these would prevent the gradients from flowing back.</p>

Chuong_Vo: <p>So even if I try to have a tensor <code>x1</code> and <code>x2</code> with gradient and replace their values with new values from the <code>preprocessXXX</code>, it is still impossible for the gradients to flow back?</p>
<p>Do you have any suggestion for this case because I had to turn the tensor to numpy and put the loop into numba to save time. If there is anyway to speed up the nested loop ut preserve the gradient please letme know! I’m kinda desperate right now <img src="https://discuss.pytorch.org/images/emoji/apple/cry.png?v=9" title=":cry:" class="emoji" alt=":cry:"> <img src="https://discuss.pytorch.org/images/emoji/apple/cry.png?v=9" title=":cry:" class="emoji" alt=":cry:"> <img src="https://discuss.pytorch.org/images/emoji/apple/cry.png?v=9" title=":cry:" class="emoji" alt=":cry:"></p>

albanD: <p>If you hide from pytorch how the values are computed, it cannot compute the gradient.</p>
<p>I’m not sure what your <code>change_patches</code> function is doing but it looks like you should be able to implement it with regular indexing functions.<br>
Could you give in detail with small example what you want from that function?</p>

Chuong_Vo: <p>So now I have the input is a tensor of size (1,1,60,60). I pad the image by 1 for every side, divide the image into 3x3 patches by using <code>unfold()</code> with stride 1, process them and then stitch them back together into a new one using <code>view()</code>. Then the new image will be 3xtimes bigger than the input image.</p>
<p>For the process, I have a dictionary having the positions according to each position of the center pixel of the 3x3 patches. And then for each patch, I will replace it’s 8 outer pixels with the pixels from the input image that have the coordinates stored in the dictionary.</p>
<pre><code class="lang-auto">x = torch.randn(1, 1, 60, 60)

coor = {(0,0) : np.matrix([[59, 7, 14, 22, 29, 36, 44, 51], [1, 1, 1, 1, 1, 1, 1, 1]])
        ...
       }
kc, kh, kw = 1, 3, 3  # kernel size
dc, dh, dw = 1, 1, 1  # stride
pad = (1,1,1,1)
paddedimage = F.pad(x, pad).to(DEVICE)
patches = paddedimage.unfold(1, kc, dc).unfold(2, kh, dh).unfold(3, kw, dw)
unfold_shape = patches.size()
print(unfold_shape)
patches = patches.contiguous().view(-1, kc, kh, kw)
print(patches.shape)

a = 0
for i in range(60):
    for j in range(60):
        coordinate = [i,j]
        indices = np.matrix(coor[(coordinate[0],coordinate[1])])
        patches[a,0,2,1] = image[0,0,indices[1,0],indices[0,0]]
        patches[a,0,2,2] = image[0,0,indices[1,1],indices[0,1]]
        patches[a,0,1,2] = image[0,0,indices[1,2],indices[0,2]]
        patches[a,0,0,2] = image[0,0,indices[1,3],indices[0,3]]
        patches[a,0,0,1] = image[0,0,indices[1,4],indices[0,4]]
        patches[a,0,0,0] = image[0,0,indices[1,5],indices[0,5]]
        patches[a,0,1,0] = image[0,0,indices[1,6],indices[0,6]]
        patches[a,0,2,0] = image[0,0,indices[1,7],indices[0,7]]
        a += 1
# Reshape back
patches_orig = patches.view(unfold_shape)
output_c = unfold_shape[1] * unfold_shape[4]
output_h = unfold_shape[2] * unfold_shape[5]
output_w = unfold_shape[3] * unfold_shape[6]
patches_orig = patches_orig.permute(0, 1, 4, 2, 5, 3, 6).contiguous()
patches_orig = patches_orig.view(1, output_c, output_h, output_w)

</code></pre>
<p>After this process, the output can become the input for the next convolutional layer…<br>
What I want is to apply this process after every conv2d() layer which means the example with 1 channel will extend to cases with many channels input.<br>
Is there anyway to optimize this process? Thank you Mr. <a class="mention" href="/u/alband">@albanD</a>!!!</p>

albanD: <p>Hi,</p>
<p>I’m confused about what <code>a</code> is doing here.<br>
The unfold shape is <code>[1, 1, 60, 60, 1, 3, 3]</code> right?<br>
So if you increment a every time, it is going to go out of bound very quickly no?</p>

Chuong_Vo: <p>I’m very sorry for the messed up code, I got mistaken <img src="https://discuss.pytorch.org/images/emoji/apple/sweat_smile.png?v=9" title=":sweat_smile:" class="emoji" alt=":sweat_smile:"><br>
I editted the code. Please check it again!<br>
I will be very grateful if you can help me with this!</p>

albanD: <p>Right.<br>
That should be significantly faster already.</p>
<p>Using fancier function and playing with views, you might be able to narrow it down to a single call. But that shouldn’t be necessary <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=9" title=":slight_smile:" class="emoji" alt=":slight_smile:"></p>
<pre><code class="lang-python">import torch
from torch.nn import functional as F
import numpy as np

image = torch.randn(1, 1, 60, 60)
coor = torch.LongTensor(60, 60, 2, 8)
coor.fill_(0) # Just to get valid indices

kc, kh, kw = 1, 3, 3  # kernel size
dc, dh, dw = 1, 1, 1  # stride
pad = (1,1,1,1)

def original(image):

    paddedimage = F.pad(image, pad)
    patches = paddedimage.unfold(1, kc, dc).unfold(2, kh, dh).unfold(3, kw, dw)
    unfold_shape = patches.size()
    patches = patches.contiguous().view(-1, kc, kh, kw)

    a = 0
    for i in range(60):
        for j in range(60):
            coordinate = [i,j]
            indices = np.matrix(coor[(coordinate[0],coordinate[1])])
            patches[a,0,2,1] = image[0,0,indices[1,0],indices[0,0]]
            patches[a,0,2,2] = image[0,0,indices[1,1],indices[0,1]]
            patches[a,0,1,2] = image[0,0,indices[1,2],indices[0,2]]
            patches[a,0,0,2] = image[0,0,indices[1,3],indices[0,3]]
            patches[a,0,0,1] = image[0,0,indices[1,4],indices[0,4]]
            patches[a,0,0,0] = image[0,0,indices[1,5],indices[0,5]]
            patches[a,0,1,0] = image[0,0,indices[1,6],indices[0,6]]
            patches[a,0,2,0] = image[0,0,indices[1,7],indices[0,7]]
            a += 1
    # Reshape back
    patches_orig = patches.view(unfold_shape)
    output_c = unfold_shape[1] * unfold_shape[4]
    output_h = unfold_shape[2] * unfold_shape[5]
    output_w = unfold_shape[3] * unfold_shape[6]
    patches_orig = patches_orig.permute(0, 1, 4, 2, 5, 3, 6).contiguous()
    patches_orig = patches_orig.view(1, output_c, output_h, output_w)

    return patches_orig

def new(image):
    paddedimage = F.pad(image, pad)
    patches = paddedimage.unfold(1, kc, dc).unfold(2, kh, dh).unfold(3, kw, dw)
    unfold_shape = patches.size()
    patches = patches.contiguous()

    coor_in = coor.unsqueeze(0).unsqueeze(0)
    patches[0, 0, :, :,0,2,1] = image[0,0,coor_in[0, 0, :, :, 1,0],coor_in[0, 0, :,:,0,0]]
    patches[0, 0, :, :,0,2,2] = image[0,0,coor_in[0, 0, :, :, 1,1],coor_in[0, 0, :,:,0,1]]
    patches[0, 0, :, :,0,1,2] = image[0,0,coor_in[0, 0, :, :, 1,2],coor_in[0, 0, :,:,0,2]]
    patches[0, 0, :, :,0,0,2] = image[0,0,coor_in[0, 0, :, :, 1,3],coor_in[0, 0, :,:,0,3]]
    patches[0, 0, :, :,0,0,1] = image[0,0,coor_in[0, 0, :, :, 1,4],coor_in[0, 0, :,:,0,4]]
    patches[0, 0, :, :,0,0,0] = image[0,0,coor_in[0, 0, :, :, 1,5],coor_in[0, 0, :,:,0,5]]
    patches[0, 0, :, :,0,1,0] = image[0,0,coor_in[0, 0, :, :, 1,6],coor_in[0, 0, :,:,0,6]]
    patches[0, 0, :, :,0,2,0] = image[0,0,coor_in[0, 0, :, :, 1,7],coor_in[0, 0, :,:,0,7]]

    # Reshape back
    patches_orig = patches.view(unfold_shape)
    output_c = unfold_shape[1] * unfold_shape[4]
    output_h = unfold_shape[2] * unfold_shape[5]
    output_w = unfold_shape[3] * unfold_shape[6]
    patches_orig = patches_orig.permute(0, 1, 4, 2, 5, 3, 6).contiguous()
    patches_orig = patches_orig.view(1, output_c, output_h, output_w)

    return patches_orig


print((original(image.clone()) - new(image.clone())).abs().max())
</code></pre>

Chuong_Vo: <p>Thank you so much Mr. <a class="mention" href="/u/alband">@albanD</a>!<br>
The code works well in a significantly faster speed!<br>
I have applied it into my training code. But the accuracy of the model seems still as low as before. The reason I think that the operations of replacing pixels by pixels may mess up the gradient of the input.<br>
Suppose that I apply this function to an output of a <code>conv2d()</code> layer, can there be a way for the output to have a decent gradient that can be backwarded to the beginning?</p>

albanD: <p>Hi,</p>
<p>No the gradients are properly computed.<br>
You can check this by running:</p>
<pre><code class="lang-python">
from torch.autograd import gradcheck
gradcheck(lambda x: new(x).sum(), image.clone().detach().double().requires_grad_())
</code></pre>
<p>It checks that the autograd gradients match the ones computed with finite difference.</p>

Chuong_Vo: <p>Dear Mr. <a class="mention" href="/u/alband">@albanD</a>,<br>
The solution that you showed me truely saved me from wasting a lot of time in my capstone project. I can’t express how grateful I am to your help and support!<br>
I wish you all the best for your career and hope if I encounter any problem in the future, I can come back and receive valuable information and suggestions from the pytorch community!</p>

cbd: <p>I need to use detach during conversion from tensor to numpy for kernel density estimation function but it want back propagate gradient. Is there any workaround solution?</p>