Hi,
I want to use a hyper-net to change the weights of a backbone-net. The following snippet is not meaningful but illustrates in a concise way what I want to do:
import sys
import torch
from torch import nn, optim
# our backbone-net
modelV = nn.Linear(1, 1, bias=False)
nn.init.zeros_(modelV.weight)
# our hyper-net
modelW = nn.Linear(1, 1, bias=False)
nn.init.constant_(modelW.weight, 0.1)
optimizer = optim.Adam(modelW.parameters(), lr=1e-3)
loss_func = nn.MSELoss()
# our data
# the hyper-net takes in z = 1 and must output 2 if trained
# the backbone-net takes in X and must output Y if trained with Y = 2 X
Z = torch.tensor(1.).reshape(-1,1)
X = torch.Tensor( [ -1. , 0 , 1., 2. ] ).reshape(-1,1)
y = torch.Tensor( [ -2., 0., 2., 4. ] ).reshape(-1,1)
for ite in range(10):
hypernet_prediction = modelW(Z)
# build a backbone-net
backbone_state_dict = modelV.state_dict()
backbone_state_dict["weight"] = hypernet_prediction[0].view(backbone_state_dict["weight"].size())
modelV.load_state_dict(backbone_state_dict)
# predict with backbone-net
backbonet_prediction = modelV(X)
# compute loss
loss = loss_func(backbonet_prediction, y)
print("ite = {}, loss = {}".format(ite, loss.item()))
# optimize
optimizer.zero_grad()
loss.backward()
optimizer.step()
This way of changing weights is very elegant but unfortunately it does not work because as mentioned in other threads the loading of the state dict is not tracked as a differentiable op by autograd.
What alternatives are there ? (keeping in mind that the real models will be much much larger)
thanks for any help