Hypernetwork implementation

mariaalfaroc: <p>Hello everyone,<br>
How can I implement this example of Hypernetworks in Pytorch? I’m having problems with the backward pass.</p><aside class="onebox allowlistedgeneric" data-onebox-src="https://keras.io/getting_started/intro_to_keras_for_researchers/#endtoend-experiment-example-2-hypernetworks">
  <header class="source">
      <img src="https://keras.io/img/favicon.ico" class="site-icon" width="64" height="64">

      <a href="https://keras.io/getting_started/intro_to_keras_for_researchers/#endtoend-experiment-example-2-hypernetworks" target="_blank" rel="noopener nofollow ugc">keras.io</a>
  </header>

  <article class="onebox-body">
    <div class="aspect-image" style="--aspect-ratio:690/239;"><img src="https://keras.io/img/logo-k-keras-wb.png" class="thumbnail" width="690" height="239"></div>

<h3><a href="https://keras.io/getting_started/intro_to_keras_for_researchers/#endtoend-experiment-example-2-hypernetworks" target="_blank" rel="noopener nofollow ugc">Keras documentation: Introduction to Keras for Researchers</a></h3>

  <p>Keras documentation</p>


  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>


InnovArul: <p>If you could post a code snippet that could reproduce the issue,  it will be easier for the community to help you.</p>

mariaalfaroc: <p>Here it’s a toy example. I want hyper_net to predict the weights for main_net. The backward pass should only take into account the hypernetwork parameters.</p>
<pre><code class="lang-auto">import torch
import torch.nn as nn
import numpy as np

main_layers = [nn.Linear(64, 8), nn.Linear(8, 8), nn.Linear(8, 1)]
main_net = nn.Sequential(*main_layers)

def get_num_weights(model):
  weights = dict()
  num_weights = 0
  for name, param in model.named_parameters():
    if param.requires_grad:
      weights[name] = tuple(param.shape)
      num_weights += torch.prod(torch.tensor(param.shape)).item()
      param.requires_grad = False
  return weights, num_weights
w, nw = get_num_weights(main_net)

hyper_layers = [nn.Linear(64, 8), nn.Linear(8, 8), nn.Linear(8, nw)]
hyper_net = nn.Sequential(*hyper_layers)

def set_weights(weights_dict, model, pred_weights):
  index = 0
  for name, param in model.named_parameters():
    if name in weights_dict.keys():
      w_coeffs = pred_weights[:, index : index + torch.prod(torch.tensor(weights_dict[name])).item()]
      w_coeffs = w_coeffs.reshape(weights_dict[name])
      param.data = w_coeffs
      index += torch.prod(torch.tensor(weights_dict[name])).item()

optimizer = torch.optim.RMSprop(hyper_net.parameters(), lr=1e-3)
criterion = nn.MSELoss()

x = torch.rand(1, 64)
y = torch.rand(1, 1)


optimizer.zero_grad()
weights = hyper_net(x)
print(weights.shape)
set_weights(weights_dict=w, model=main_net, pred_weights=weights)
yhat = main_net(x)
loss = criterion(yhat, y)
loss.backward()
optimizer.step()
</code></pre>
<p>I get the following error:</p>
<blockquote>
<p>RuntimeError: element 0 of tensors does not require grad and does not have a grad_fn</p>
</blockquote>
<p>If I don’t set requires_grad to False, it works fine. However, I don’t know if it is working as I want. The main network should not take part into the gradient computation, I only use to make predictions based on a set of predicted weights. It’s basically a frozen model. In keras, they use</p>
<pre><code class="lang-auto">for layer in main_network.layers:
    layer.built = True
</code></pre>
<p>However, I don’t know if we have any equivalent for the built attribute in Pytorch.</p>

mariaalfaroc: <p>You are right! Sorry. Its my first time posting here. Also, please forgive my english!</p>

ID56: <aside class="quote no-group" data-username="mariaalfaroc" data-post="3" data-topic="152109">
<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/mariaalfaroc/48/49777_2.png" class="avatar"> mariaalfaroc:</div>
<blockquote>
<p><code>param.requires_grad = False</code></p>
</blockquote>
</aside>
<p>No need to turn grad off for the main model, since the optimizer will update only the hypernet’s params. If you turn off gradients, your prediction / loss will also lack gradients, and thus fail.</p>
<aside class="quote no-group" data-username="mariaalfaroc" data-post="3" data-topic="152109">
<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/mariaalfaroc/48/49777_2.png" class="avatar"> mariaalfaroc:</div>
<blockquote>
<p><code>param.data = w_coeffs</code></p>
</blockquote>
</aside>
<p>Here is the second issue, and the main caveat: you can’t simply assign predictions of one model as a weight parameter of another model. Well you can, but that would detach model 1 (hypernetwork) from the computation graph. Consider the following:</p>
<pre data-code-wrap="python"><code class="lang-python">import torch
from torch import nn, optim

A = nn.Linear(1, 1) # main network
B = nn.Linear(1, 2) # hypernetwork &lt;-- need to optimize this

def foo(a, b, c):
    print("B Forward")

def bar(a, b, c):
    print("B Backward")

B.register_forward_hook(foo)
B.register_full_backward_hook(bar)

optimizer = optim.Adam(B.parameters(), lr=1e-3)

X = torch.randn(1)

optimizer.zero_grad()

W, b = B(X)
A.weight.data = W.reshape(1, 1)
B.bias.data = b

pred = A(X)
loss = pred.sum()
loss.backward()
optimizer.step()
</code></pre>
<p>If we run this we will get:</p>
<pre><code class="lang-auto">B Forward
</code></pre>
<p>“B Backward” isn’t printing. B isn’t getting any gradients!</p>
<p>The workaround for this could be by simply using the backend <code>nn.functional</code> layers. Instead of <code>nn.Linear()(x)</code>, we will have to use <code>F.linear(x, W, b)</code> to retain the gradients.</p>
<p>It is pretty simple to just convert the toy keras example you linked with only a few changes:</p>
<pre data-code-wrap="python"><code class="lang-python">import torch
from torch import nn, optim
import torch.nn.functional as F
import numpy as np
import tensorflow as tf

def main_network(x, params):
    w0, b0 = params[0]
    w1, b1 = params[1]

    x = F.relu(F.linear(x, w0, b0))
    x = F.linear(x, w1, b1)
    return x

input_dim = 784
classes = 10

layer_shapes = [
    [(64, input_dim), (64,)],   # (w0, b0)
    [(classes, 64), (classes,)] # (w1, b1)
]

num_weights_to_generate = (classes * 64 + classes) + (64 * input_dim + 64)

hypernetwork = nn.Sequential(
    nn.Linear(input_dim, 16),
    nn.ReLU(),
    nn.Linear(16, num_weights_to_generate),
    nn.Sigmoid()
)

loss_fn = nn.CrossEntropyLoss()
optimizer = optim.Adam(hypernetwork.parameters(), lr=1e-4)

def train_step(x, y):
    optimizer.zero_grad()
    weights_pred = hypernetwork(x)
 
    idx, params = 0, []
    for layer in layer_shapes:
        layer_params = []
        for shape in layer:
            offset = np.prod(shape)
            layer_params.append(weights_pred[:, idx: idx + offset].reshape(shape))
            idx += offset
        params.append(layer_params)

    preds = main_network(x, params)
    loss = loss_fn(preds, y)
    loss.backward()
    optimizer.step()
    return loss.item()
</code></pre>
<p>You can test it with keras’s dataset too, just convert those tf tensors to torch tensors via numpy.</p>
<pre data-code-wrap="python"><code class="lang-python">(x_train, y_train), _ = tf.keras.datasets.mnist.load_data()
dataset = tf.data.Dataset.from_tensor_slices(
    (x_train.reshape(60000, 784).astype("float32") / 255, y_train)
)
dataset = dataset.shuffle(buffer_size=1024).batch(1)

losses = []

torchify = lambda a: torch.from_numpy(a.numpy())

for step, (x, y) in enumerate(dataset):
    x, y = torchify(x).float(), torchify(y).long()
    loss = train_step(x, y)
    losses.append(loss)
    if step % 1 == 0: print("Step:", step, "Loss:", sum(losses) / len(losses))
    if step &gt;= 1000: break
</code></pre>
<p>Hope this example helps!</p>

mariaalfaroc: <p>This example is amazing! You absolutely solve my problem.</p>
<p>Sorry for the late reply! Work and life got in the way.<br>
I have just replicated your answer and also tried it on a bigger model, and everything works perfectly! I cannot thank you enough <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=12" title=":slight_smile:" class="emoji" alt=":slight_smile:" loading="lazy" width="20" height="20"></p>

mariaalfaroc: <p>Hello!<br>
I have some doubts…what if I want to do a hypernetwork that predicts the weights for a LSTM layer? I’ve been looking for the functional version of LSTM, but there isn’t one. Should I use this one?</p><aside class="onebox githubblob" data-onebox-src="https://github.com/pytorch/pytorch/blob/master/torch/nn/modules/rnn.py">
  <header class="source">

      <a href="https://github.com/pytorch/pytorch/blob/master/torch/nn/modules/rnn.py" target="_blank" rel="noopener nofollow ugc">github.com</a>
  </header>

  <article class="onebox-body">
    <h4><a href="https://github.com/pytorch/pytorch/blob/master/torch/nn/modules/rnn.py" target="_blank" rel="noopener nofollow ugc">pytorch/pytorch/blob/master/torch/nn/modules/rnn.py</a></h4>


      <pre><code class="lang-py">import math
import warnings
import numbers
from typing import List, Tuple, Optional, overload

import torch
from torch import Tensor
from .module import Module
from ..parameter import Parameter
from ..utils.rnn import PackedSequence
from .. import init
from ... import _VF

__all__ = ['RNNBase', 'RNN', 'LSTM', 'GRU', 'RNNCellBase', 'RNNCell', 'LSTMCell', 'GRUCell']

_rnn_impls = {
    'RNN_TANH': _VF.rnn_tanh,
    'RNN_RELU': _VF.rnn_relu,
}

</code></pre>



  This file has been truncated. <a href="https://github.com/pytorch/pytorch/blob/master/torch/nn/modules/rnn.py" target="_blank" rel="noopener nofollow ugc">show original</a>

  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>

<p>from torch import _VF<br>
_VF.lstm()</p>

mariaalfaroc: <p><a class="mention" href="/u/id56">@ID56</a> I’ve been thinking about this again. I was wondering if then, using <code>torch.nn.Functional</code> is the only way to implement hypernetworks. Is there a way to make it work using <code>torch.nn.Parameter</code>? If I try to do so using your example:</p>
<pre><code class="lang-auto">import torch
from torch import nn, optim
import torch.nn.functional as F
import numpy as np
import tensorflow as tf

input_dim = 784
classes = 10

layer_shapes = [
    [(64, input_dim), (64,)],   # (w0, b0)
    [(classes, 64), (classes,)] # (w1, b1)
]

num_weights_to_generate = (classes * 64 + classes) + (64 * input_dim + 64)

hypernetwork = nn.Sequential(
    nn.Linear(input_dim, 16),
    nn.ReLU(),
    nn.Linear(16, num_weights_to_generate),
    nn.Sigmoid()
)

loss_fn = nn.CrossEntropyLoss()
optimizer = optim.Adam(hypernetwork.parameters(), lr=1e-4)

class MainNet(torch.nn.Module):
    def __init__(self):
        super(MainNet, self).__init__()
        self.linear1 = torch.nn.Linear(input_dim, 64)
        self.relu = torch.nn.ReLU()
        self.linear2 = torch.nn.Linear(64, classes)

    def forward(self, x):
        with torch.no_grad():
            x = self.linear1(x)
            x = self.relu(x)
            x = self.linear2(x)
            return x

    def set_weights(self, params):
        with torch.no_grad():
            w0, b0 = params[0]
            w1, b1 = params[1]
            self.linear1.weight = torch.nn.Parameter(w0, requires_grad=False)
            self.linear1.bias = torch.nn.Parameter(b0, requires_grad=False)
            self.linear2.weight = torch.nn.Parameter(w1, requires_grad=False)
            self.linear2.bias = torch.nn.Parameter(b1, requires_grad=False)

main_network = MainNet()

def train_step(x, y):
    optimizer.zero_grad()
    weights_pred = hypernetwork(x)
 
    idx, params = 0, []
    for layer in layer_shapes:
        layer_params = []
        for shape in layer:
            offset = np.prod(shape)
            layer_params.append(weights_pred[:, idx: idx + offset].reshape(shape))
            idx += offset
        params.append(layer_params)

    main_network.set_weights(params)
    preds = main_network(x)
    loss = loss_fn(preds, y)
    loss.backward(retain_graph=True)
    optimizer.step()
    return loss.item()

(x_train, y_train), _ = tf.keras.datasets.mnist.load_data()
dataset = tf.data.Dataset.from_tensor_slices(
    (x_train.reshape(60000, 784).astype("float32") / 255, y_train)
)
dataset = dataset.shuffle(buffer_size=1024).batch(1)

losses = []

torchify = lambda a: torch.from_numpy(a.numpy())

for step, (x, y) in enumerate(dataset):
    x, y = torchify(x).float(), torchify(y).long()
    loss = train_step(x, y)
    losses.append(loss)
    if step % 1 == 0: print("Step:", step, "Loss:", sum(losses) / len(losses))
    if step &gt;= 1000: break
</code></pre>
<p>If I remove all the <code>torch.no_grad()</code> and set <code>requires_grad = True</code>, it works. But I’m worried that I’m doing something wrong, because the gradients for MainNetwork are also being computed. Even though I think this is not a problem since those new weights that backprop computes are never used since they are overwrite by Hypernetwork. So, the only problem would be that it would take longer to optimize the Hypernetwork.</p>

mariaalfaroc: <p>I’m back again hahah In Pytorch 1.12.0, there exists <code>torch.nn.utils.stateless.functional_call</code>, which I recently discovered thanks to <a class="mention" href="/u/alband">@albanD</a> <img src="https://discuss.pytorch.org/images/emoji/apple/grin.png?v=12" title=":grin:" class="emoji" alt=":grin:" loading="lazy" width="20" height="20"> (see: <a href="https://github.com/pytorch/pytorch/issues/80454#issuecomment-1175244653" class="inline-onebox" rel="noopener nofollow ugc">Functional version of torch.nn.LSTM -&gt; torch.nn.functional.LSTM · Issue #80454 · pytorch/pytorch · GitHub</a>).</p>
<p>That way we can create the main network as another module but call it in a functional way. This is great for using, for example, LSTM which do not have a functional layer per se. So the above example could be solved in the following way:</p>
<pre><code class="lang-auto">import torch
from torch import nn, optim
import torch.nn.functional as F
import numpy as np
import tensorflow as tf

input_dim = 784
classes = 10

layer_shapes = [
    [(64, input_dim), (64,)],   # (w0, b0)
    [(classes, 64), (classes,)] # (w1, b1)
]

num_weights_to_generate = (classes * 64 + classes) + (64 * input_dim + 64)

hypernetwork = nn.Sequential(
    nn.Linear(input_dim, 16),
    nn.ReLU(),
    nn.Linear(16, num_weights_to_generate),
    nn.Sigmoid()
)

loss_fn = nn.CrossEntropyLoss()
optimizer = optim.Adam(hypernetwork.parameters(), lr=1e-4)

class MainNet(torch.nn.Module):
    def __init__(self):
        super(MainNet, self).__init__()
        self.linear1 = torch.nn.Linear(input_dim, 64)
        self.relu = torch.nn.ReLU()
        self.linear2 = torch.nn.Linear(64, classes)

    def forward(self, x):
        x = self.linear1(x)
        x = self.relu(x)
        x = self.linear2(x)
        return x

main_network = MainNet()
for param in main_network.parameters():
    param.requires_grad = False

def format_weights(params):
    w0, b0 = params[0]
    w1, b1 = params[1]
    weights = {
        "linear1.weight": w0,
        "linear1.bias": b0,
        "linear2.weight": w1,
        "linear2.bias": b1,
    }
    return weights

def train_step(x, y):
    optimizer.zero_grad()
    weights_pred = hypernetwork(x)
 
    idx, params = 0, []
    for layer in layer_shapes:
        layer_params = []
        for shape in layer:
            offset = np.prod(shape)
            layer_params.append(weights_pred[:, idx: idx + offset].reshape(shape))
            idx += offset
        params.append(layer_params)

    weights = format_weights(params)
    preds = torch.nn.utils.stateless.functional_call(main_network, weights, x)
    loss = loss_fn(preds, y)
    loss.backward(retain_graph=True)
    optimizer.step()
    return loss.item()



(x_train, y_train), _ = tf.keras.datasets.mnist.load_data()
dataset = tf.data.Dataset.from_tensor_slices(
    (x_train.reshape(60000, 784).astype("float32") / 255, y_train)
)
dataset = dataset.shuffle(buffer_size=1024).batch(1)

losses = []

torchify = lambda a: torch.from_numpy(a.numpy())

for step, (x, y) in enumerate(dataset):
    x, y = torchify(x).float(), torchify(y).long()
    loss, preds = train_step(x, y)
    losses.append(loss)
    if step % 1 == 0: print("Step:", step, "Loss:", sum(losses) / len(losses))
    if step &gt;= 1000: break
</code></pre>
<p>Hope this helps!</p>

mattvillain: <p>Now that it has been deprecated, use</p>
<p>torch.func.functional_call instead of torch.nn.utils.stateless.functional_call, otherwise should be the same.</p>
<p><a href="https://pytorch.org/docs/stable/generated/torch.func.functional_call.html#torch.func.functional_call" class="onebox" target="_blank" rel="noopener nofollow ugc">https://pytorch.org/docs/stable/generated/torch.func.functional_call.html#torch.func.functional_call</a></p>
<p>I am still not super clear on whether we should use ‘require_grad = False’ or not if we specify which parameters to optimise over in the optimiser class.</p>
<p>Have a good one!</p>
<p><img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=12" title=":slight_smile:" class="emoji only-emoji" alt=":slight_smile:" loading="lazy" width="20" height="20"></p>