Optimizing regression weights for NN outputs with PyTorch

Question

So I'm basically trying to fit a regression on the relation of the input and output of a neural network model. Then the idea is, that these estimated regression weights should be optimized to some specific target value (let's say matrix of ones). So the weights are optimized, but have a direct relation to the neural network weights. For a simple problem, I can find a solution on paper, but I fail to implement this. The loss is not decreasing, but oscillating.

# Some arbitrary PyTorch model
my_nn = Net()

# Our target regression weights
target = torch.ones(2, 2)

# Optimizer for NN weights
optimizer = optim.Adam(my_nn.parameters(), lr=0.0001)

# Arbitrary loss function
criterion = nn.MSELoss()

# Training loop
my_nn.train()
for epoch in range(500):
    print('Running Epoch: ', epoch)
    # zero the parameter gradients
    optimizer.zero_grad()
    
    # output of model is 2-dim
    y_pred = my_nn(X)

    # weights for regression model
    w1 = Variable(torch.randn(2, 2).type(torch.float), requires_grad=True)
    # optimizer for regression model
    optimizer1 = optim.SGD([w1], lr=0.001, momentum=0.9)
    for nest_epoch in range(1000):
        optimizer1.zero_grad()
        
        y_hat = torch.mm(X, w1)
        loss = criterion(y_hat, y_pred)

        # Need to add retain_graph, as it doesn't work otherwise
        # but it does work, this loss will get close to 0
        loss.backward(retain_graph=True)
        optimizer1.step()

    # this loss does not decrease
    loss = criterion(target, w1)
    print('Loss:', loss.item())
    loss.backward()
    optimizer.step()
```

Answer 1

My solution is to update the variable manually and creating a new graph for autograd. Works perfectly well now!

my_nn.train()
for epoch in range(300):
    # zero the parameter gradients
    optimizer.zero_grad()
    my_nn.zero_grad()

    Y = my_nn(X)
    loss_output = criterion(Y, y)
    w = torch.zeros_like(target)
    w.requires_grad = True
    for nest_epoch in range(50):
        y_w = torch.matmul(X, w)
        loss_w = criterion(y_w, Y)
        grad = torch.autograd.grad(loss_w, w, create_graph=True)
        w = w - 0.1 * grad[0]
    loss_regression = criterion(w, target)
    loss = loss_output + loss_regression
    loss.backward()
    optimizer.step()