Confusion about nested modules and shared parameters

snakeztc: <p>I am confused about the relationship between assign a module as another module’s class member vs. reused the module. Supposed I want to share word embedding between two networks, I used to write code as the following:</p>
<pre><code class="lang-auto">
class Model(nn.Module):
    def __init__():
        self.embedding = nn.Embedding(10000, 200)
        self.net_a = SubModule(self.embedding)
        self.net_b = SubModule(self.embedding)

   def forward(input):
       return self.net_a(input) + self.net_b(input)

class SubModule(nn.Module):
    def __init___(embedding):
        self.embedding = embedding
        self.fc = nn.Linear(200, 200)

   def forward(self, input):
        return self.fc(self.embedding(input))

</code></pre>
<p>In the above code, is the embedding correctly shared between net_a and net_b? It seems that the constructor of SubModule create another copy of the embedding so that there are three embedding matrix stored in the model?</p>

InnovArul: <p>Yes. Its a correct way of sharing the module.</p>
<p>After correcting some obvious bugs in your code (super class <code>__init__()</code>, indentation, etc), below is the working code. You can use <code>model.named_parameters()</code> to check the learnable parameters:</p>
<pre><code class="lang-auto">import torch
import torch.nn as nn


class SubModule(nn.Module):
    def __init__(self, embedding):
        super(SubModule, self).__init__()
        self.embedding = embedding
        self.fc = nn.Linear(200, 200)

    def forward(self, input):
        return self.fc(self.embedding(input))


class Model(nn.Module):
    def __init__(self):
        super(Model,self).__init__()
        self.embedding = nn.Embedding(10000, 200)
        self.net_a = SubModule(self.embedding)
        self.net_b = SubModule(self.embedding)

    def forward(self, input):
        return self.net_a(input) + self.net_b(input)

m = Model()
for n, p in m.named_parameters():
    print(n, p.shape)

</code></pre>
<p>output:</p>
<pre><code class="lang-auto">('embedding.weight', (10000, 200))
('net_a.fc.weight', (200, 200))
('net_a.fc.bias', (200,))
('net_b.fc.weight', (200, 200))
('net_b.fc.bias', (200,))

</code></pre>

snakeztc: <p>Thanks! However, if you print the keys in m.state_dict(), you will find:</p>
<pre><code class="lang-auto">['embedding.weight',
 'net_a.embedding.weight',
 'net_a.fc.weight',
 'net_a.fc.bias',
 'net_b.embedding.weight',
 'net_b.fc.weight',
 'net_b.fc.bias']
</code></pre>
<p>Are all the embedding weights above tied?</p>

InnovArul: <p>Good point!<br>
I just verified that it is handled in pytorch to still share the weights such that changing the <code>embedding</code> weight in one submodule (eg., <code>net_a</code>) affects the main <code>embedding</code> weight and other submodule (eg., <code>net_b</code>).</p>
<aside class="onebox githubgist">
  <header class="source">
      <a href="https://gist.github.com/InnovArul/d66a9e2bf56507dd0803135fb2d910d4" target="_blank">gist.github.com</a>
  </header>
  <article class="onebox-body">
    <h4><a href="https://gist.github.com/InnovArul/d66a9e2bf56507dd0803135fb2d910d4" target="_blank">https://gist.github.com/InnovArul/d66a9e2bf56507dd0803135fb2d910d4</a></h4>
<h5>pytorch_shared_module_weights.py</h5>
<pre><code class="Python">import torch
import torch.nn as nn


class SubModule(nn.Module):
    def __init__(self, embedding):
        super(SubModule, self).__init__()
        self.embedding = embedding
        self.fc = nn.Linear(200, 200)
</code></pre>
This file has been truncated. <a href="https://gist.github.com/InnovArul/d66a9e2bf56507dd0803135fb2d910d4" target="_blank">show original</a>

<p>
</p>

  </article>
  <div class="onebox-metadata">
    
    
  </div>
  <div style="clear: both"></div>
</aside>


Anubhab_Ghosh: <p>Hi, nice idea.<br>
I also had a similar question, like this kind of embedding sharing between networks also applies to share <code>nn.Sequential</code> modules? Like can <code>nn.Sequential( )</code> based ‘<code>lambdas</code>’ (with their weights inside <code>nn.Sequential()</code>) be shared across different objects of a class</p>
<p>For e.g. instead of <code>self.embedding</code>, what if I wanted something like:</p>
<pre><code class="lang-auto">import torch
import torch.nn as nn

class SubModule(nn.Module): 

    def __init__(self, net_s):
        super(SubModule, self).__init__()
        self.nets = lambda : nn.Sequential(nn.Linear(100,200), nn.ReLU())

    def forward(self, input):
        return self.nets(input)

class Model(nn.Module):
   
    def __init__(self, x, net_s):
        super(Model,self).__init__()
        self.nets1 = SubModule(net_s)
        self.nets2 = SubModule(net_s)`
    def forward(self, input):
        return self.net_a(input) + self.net_b(input)


</code></pre>
<p>Do nets1 and nets2 use the same sequential lambda defined in SubModule (with the same weights)??</p>
<p>Or am I missing something?</p>