Optimize training data after training stage

lovettxh: <p>Hi, I’m trying to optimize a noise which is used during model training to aim for a specific target when the training is done.<br>
A sample code can be found here:</p>
<pre><code class="lang-auto">import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision
from torchvision import transforms

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(1, 32, 3, 1)
        self.conv2 = nn.Conv2d(32, 64, 3, 1)
        self.dropout1 = nn.Dropout(0.25)
        self.dropout2 = nn.Dropout(0.5)
        self.fc1 = nn.Linear(9216, 128)
        self.fc2 = nn.Linear(128, 10)

    def forward(self, x):
        x = self.conv1(x)
        x = F.relu(x)
        x = self.conv2(x)
        x = F.relu(x)
        x = F.max_pool2d(x, 2)
        x = self.dropout1(x)
        x = torch.flatten(x, 1)
        x = self.fc1(x)
        x = F.relu(x)
        x = self.dropout2(x)
        x = self.fc2(x)
        output = F.log_softmax(x, dim=1)
        return output

def train(model, optimizer, data, label, noise, target):

    noise.requires_grad_()
    # model training
    model.train()
    optimizer.zero_grad()
    output = model(data + noise)
    loss = F.cross_entropy(output, label)
    noise.retain_grad()
    loss.backward()
    optimizer.step()
    
    # optimize noise
    model.eval()
    output = model(target['data'])
    loss = -F.cross_entropy(output, target['label'])
    grad = torch.autograd.grad(loss, [noise])[0]
    noise = noise.detach() + 0.001 * torch.sign(grad.detach())
    noise = torch.clamp(noise, 0, 1)
    return noise

def main():
    model = Net().cuda()
    optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
    data = torch.randn((100, 1, 28, 28)).cuda()
    label = torch.randint(0, 10, (100,)).cuda()

    noise = torch.randn((100, 1, 28, 28)).cuda()
    target = {'data': torch.randn((1, 1, 28, 28)).cuda(), 'label': torch.tensor([0]).cuda()}
    train(model, optimizer, data, label, noise, target)

if __name__ == '__main__':
    main()
</code></pre>
<p>When I run the code it says <code>RuntimeError: One of the differentiated Tensors appears to not have been used in the graph. Set allow_unused=True if this is the desired behavior.</code> on the line <code>grad = torch.autograd.grad(loss, [noise])[0]</code>.</p>
<p>If I add allow_unused=True to it, the grad will just be None. It seems that the gradient of the noise wrt the model during training is gone when I optimize the noise.</p>
<p>I wonder if there is any way to retain the gradient of the noise so I can optimize it later.</p>
<p>I also tried <code>loss.backward(create_graph=True)</code>, but it didn’t work.</p>

ptrblck: <p>The <code>noise</code> tensor wasn’t used in the eval loop in the computation of <code>output</code> and <code>loss</code> so it also won’t influence their values resulting in a <code>None</code> gradient. You would need to add <code>noise</code> to the actual computation so that its value influences the output.</p>

lovettxh: <p>Hi, I understand that the <code>noise</code> tensor is not directly involved in the loss calculation in the evaluation part. But it is involved in the model training part, so the model parameters in evaluation should be calculated through <code>noise</code>, which I guess makes it indirectly involve in the loss computation.</p>

lovettxh: <p>Hi, I managed to work on something simpler with these codes.</p>
<pre><code class="lang-auto">import torch
import torch.nn.functional as F
a = torch.randn(2,2,requires_grad=True)
b = torch.randn(2,2)
noise = torch.randn(2,2,requires_grad=True)
loss1 = torch.mm(a,b+noise)
loss1 = F.cross_entropy(loss1, torch.tensor([0,1]))
# grad1 = torch.autograd.grad(loss1,a,grad_outputs=torch.ones_like(loss1), create_graph=True)[0]
loss1.backward(create_graph=True)

a = a - a.grad * 0.1

c = torch.randn(2,2,requires_grad=True)

loss2 = torch.mm(a,c)
loss2 = F.cross_entropy(loss2, torch.tensor([1,0]))
grad2 = torch.autograd.grad(loss2, noise, grad_outputs=torch.ones_like(loss2))[0]
print("grad2: ", grad2)
</code></pre>
<p>I think they are similar with <code>a</code> being the model, <code>b</code> being data, and <code>c</code> is the <code>target['data']</code>. However, this code works perfectly with no error. I don’t quite understand where I’m getting wrong with the first code. Please point out if I’m doing anything wrong. Thanks!</p>

lovettxh: <p>I did some research and found that the model weights update can only be done with <code>torch.no_grad()</code>.<br>
One other way to do it is through weight assignment <code>model.fc1.weight = torch.nn.Parameter(model.fc1.weight - model.fc1.weight.grad)</code>, however doing so will remove its gradient.<br>
i.e.</p>
<pre><code class="lang-auto">model.fc1.weight = torch.nn.Parameter(model.fc1.weight - model.fc1.weight.grad)
print(model.fc1.weight.grad) # --&gt; None
</code></pre>
<p>While the model is a leaf Variable and requires gradient, I can’t also use in-place operation to update weights.</p>
<p>I’m wondering if there is no solution to do this in pytorch, or if I’m still missing something.<br>
Thanks!</p>