Train a model to output weights of another model, and use the other model just as function evaluation

Ritam_9: <p>I have 2 models, A and B.<br>
A(x1)=Weights of B<br>
B(x2)=Final output</p>
<p>A is trainable<br>
B is not trainable (I just want to upload the outputs of A into B and infer)</p>
<p>Problem I am facing: Output of A is torch.tensor. While setting the weights of B, I had to slice the output tensor of A. However, I am losing the gradient flow, from final loss to weights of A, hence there is no training happening. How do I implement the idea or correct my code?</p>
<p>My Source-Code:</p>
<pre><code class="lang-auto">import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.nn.parameter import Parameter
from torch.autograd import Variable
import numpy as np

class Hyper_Model(nn.Module):

    def __init__(self):

        super(Hyper_Model, self).__init__()
        self.layers = nn.Sequential(nn.Linear(1,32),
                      nn.ReLU(),             
                      nn.Linear(32,32),
                      nn.ReLU(),
                      nn.Linear(32,32),
                      nn.ReLU(),
                      nn.Linear(32,32),
                      nn.ReLU(),
                      nn.Linear(32,177))

  
    def forward(self,param):        
        param_ = self.layers(param)            
        return param_`Preformatted text`

class Main_Model(nn.Module):

    def __init__(self):

        super(Main_Model, self).__init__()
        self.linear1 = nn.Linear(2,8)
        self.linear2 = nn.Linear(8,8)
        self.linear3 = nn.Linear(8,8)
        self.out = nn.Linear(8,1)

    def forward(self,param_,x):
        self.linear1.weight = torch.nn.Parameter(param_[0,:16].view(8,2))
        self.linear2.weight = torch.nn.Parameter(param_[0,24:88].view(8,8))
        self.linear3.weight = torch.nn.Parameter(param_[0,96:160].view(8,8))
        self.linear1.bias = torch.nn.Parameter(param_[0,16:24].view(8))
        self.linear2.bias = torch.nn.Parameter(param_[0,88:96].view(8))
        self.linear3.bias = torch.nn.Parameter(param_[0,160:168].view(8))
        self.out.weight = torch.nn.Parameter(param_[0,168:176].view(1,8))
        self.out.bias = torch.nn.Parameter(param_[0,176:].view(1))

        self.linear1.weight.requires_grad = False
        self.linear2.weight.requires_grad = False
        self.linear3.weight.requires_grad = False        
        self.linear1.bias.requires_grad = False
        self.linear2.bias.requires_grad = False
        self.linear3.bias.requires_grad = False
        self.out.weight.requires_grad =  False
        self.out.bias.requires_grad =  False

        x = F.relu(self.linear1(x))
        x = F.relu(self.linear2(x))
        x = F.relu(self.linear3(x)) 
        x = self.out(x)
        return x

x = torch.tensor([1.0,2.0,3.0],requires_grad=True).view(3,1)
t = torch.tensor([1.0,1.5,2.0],requires_grad=True).view(3,1)
param = torch.tensor([-0.01]).view(1,1)
X = torch.cat([x,t],dim=1)
Y = torch.tensor([5.0,6.0,9.0]).view(3,1)
h = Hyper_Model()
m = Main_Model()
opt = torch.optim.Adam(list(h.parameters()), lr=0.001)
loss_func = nn.MSELoss()

for i in range(10):
    opt.zero_grad()
    param_ = h(param)   

    out = m(param_,X)
    loss = loss_func(out,Y)

    print(i,loss)

    loss.backward()
    opt.step()
</code></pre>
<p>Output:<br>
0 tensor(46.0043, grad_fn=)<br>
1 tensor(46.0043, grad_fn=)<br>
2 tensor(46.0043, grad_fn=)<br>
3 tensor(46.0043, grad_fn=)<br>
4 tensor(46.0043, grad_fn=)<br>
5 tensor(46.0043, grad_fn=)<br>
6 tensor(46.0043, grad_fn=)<br>
7 tensor(46.0043, grad_fn=)<br>
8 tensor(46.0043, grad_fn=)<br>
9 tensor(46.0043, grad_fn=)</p>

Dazitu616: <p>Maybe helpful. I recently encounter a repo which does similar thing (predict the weight of another network) <a href="https://github.com/yilundu/gem" rel="noopener nofollow ugc">here</a>. <a href="https://github.com/yilundu/gem/blob/f50d279a4feec5c6af861cd36d14d934900d57b9/modules.py#L42" rel="noopener nofollow ugc">This</a> is how they do for a FL layer, inheriting from <code>MetaModule</code> from the package <code>torchmeta</code></p>

KFrank: <p>Hi Ritam!</p>
<aside class="quote no-group" data-username="Ritam_9" data-post="1" data-topic="145678" data-full="true">
<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/r/3ec8ea/48.png" class="avatar"> Ritam_9:</div>
<blockquote>
<p>A is trainable<br>
B is not trainable (I just want to upload the outputs of A into B and infer)</p>
</blockquote>
</aside>
<p>As I understand it, you do not want to directly train the weights of B, but<br>
you do want to train the weights of A so that it produces values that work<br>
better when used as the weights of B.</p>
<aside class="quote no-group">
<blockquote>
<pre><code class="lang-auto">from torch.autograd import Variable
</code></pre>
</blockquote>
</aside>
<p>(As an aside, don’t use <code>torch.autograd.Variable</code>.  It is deprecated.<br>
Just use regular pytorch tensors.)</p>
<aside class="quote no-group">
<blockquote>
<pre><code class="lang-auto">class Main_Model(nn.Module):

    def __init__(self):

        super(Main_Model, self).__init__()
        self.linear1 = nn.Linear(2,8)
...

    def forward(self,param_,x):
        self.linear1.weight = torch.nn.Parameter(param_[0,:16].view(8,2))
...
        x = F.relu(self.linear1(x))
...
        return x
</code></pre>
</blockquote>
</aside>
<p>Don’t create instances of <code>Linear</code> in <code>Main_Model</code> (as I understand it,<br>
“Model B”).  (So you won’t be setting the weights of instances of <code>Linear</code><br>
to the weight-values produced by “Model A” and you won’t be applying<br>
those instances of <code>Linear</code> to <code>x</code>.)</p>
<p>Instead, apply the non-class, <a href="https://pytorch.org/docs/1.10.0/generated/torch.nn.functional.linear.html" rel="noopener nofollow ugc">functional form</a> of <code>Linear</code> to <code>x</code> passing in<br>
the weight-values produces by “Model A” as arguments:</p>
<pre data-code-wrap="python"><code class="lang-python">    def forward (self, param_, x):
        x = F.relu (F.linear (x, weight = param_[0, :16].view (8, 2), bias = param_[0, 16:24].view (8))
        ...
        return x
</code></pre>
<p><code>F.Linear</code> is now just an ordinary pytorch tensor function.  The gradients<br>
produced by calling <code>.backward()</code> on the loss calculated from the output<br>
of “Model B” will backpropagate properly through the calls to <code>F.linear()</code><br>
in “Model B”'s <code>forward()</code> function and then back propagate through the<br>
“Model A” that produced the weight-values, producing gradients for the<br>
weights in “Model A” which will then be optimized when you call <code>opt.step()</code>.</p>
<p>Best.</p>
<p>K. Frank</p>

Ritam_9: <p>Thank you very much. I am working on something very similar to the idea in the Repo you just shared. It was very helpful!</p>

Ritam_9: <p><a class="mention" href="/u/kfrank">@KFrank</a>  Thank you very much. This is precisely what I was looking for, and is working very smoothly for my use case! Just a small follow-up, using .nn modules for proper backprop training is recommended, and .nn.functional modules are suitable for function evaluations without gradients. Will that be a crude but suitable summary?</p>

frozencandy: <p>What are the benefits of using a generator for weight generation? I tried searching the web, but didn’t get any satisfactory answer.</p>