Python dictionary in the model not trained?

hongtaesuk: <p>Hey guys,<br>
can anyone of you could explain why such things below happened?</p>
<pre><code class="lang-auto"> def __init__(self, config):
        super(BertForMultitask, self).__init__(config)
        self.bert = BertModel(config)
        self.dropout = nn.Dropout(config.hidden_dropout_prob)
        self.apply(self.init_weights)
        self.multitask_ratio = self.set_update_ratio(config, 0.7)

        #####################
        ####    Part A. This works     ####
        #####################
        self.classifier_student = nn.Linear(config.hidden_size, 2)
        self.classifier_squad = nn.Linear(config.hidden_size, 2)
        self.classifier_usermatch = nn.Linear(config.hidden_size, 2)
        self.classifier_stsb = nn.Linear(config.hidden_size, 1)

        ########################
        ####    Part B. Below doesn't     ####
        ########################        
        self.classifier = {}
        for task in config.tasks:
             if task in ['student_response', 'squad', 'user_match']:
                 self.classifier[task] = nn.Linear(config.hidden_size, self.config.num_labels)
             elif task in ['sts-b']:
                 self.classifier[task] = nn.Linear(config.hidden_size, 1)
</code></pre>
<p>I’ve been implementing multitask model that deals with several different tasks<br>
and I’ve trained with each of different Linear layer for each task after the inputs<br>
pass through a BERT model.</p>
<p>However, the problem here is that when I defined the Linear layers for each task<br>
just like the part B - dictionary holding each linear layers for each tasks - then it seems<br>
the model doesn’t hold the trained parameters of that layer after it completes the training.<br>
Part A where each layer is defined separately, not aggregated in a single python dictionary,<br>
works just fine.</p>
<p>Is there any secret behind this means of defining layers?<br>
Anyone knows any clue?</p>

ptrblck: <p>Plain Python containers, such as <code>list</code> and <code>dict</code> won’t be properly registered inside your module, so use <a href="https://pytorch.org/docs/stable/nn.html#torch.nn.ModuleDict">nn.ModuleDict</a> in that case (or <code>nn.ModuleList</code> instead of <code>list</code>).</p>

hongtaesuk: <p>Thanks a lot for the clear reply!!<br>
Should I just fix the normal dictionary like</p>
<pre><code class="lang-auto">self.classifier = nn.ModuleDict({})
</code></pre>
<p>then add elements??</p>
<p>And what do you mean by ‘not properly registered inside the module’??<br>
Does that mean that the model won’t be able to properly record operations and data flows<br>
for that variable??</p>

ptrblck: <p>Just create it via:</p>
<pre data-code-wrap="python"><code class="lang-python">self.classifier = nn.ModuleDict()
</code></pre>
<aside class="quote no-group" data-username="hongtaesuk" data-post="3" data-topic="64498">
<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/hongtaesuk/48/18907_2.png" class="avatar"> hongtaesuk:</div>
<blockquote>
<p>Does that mean that the model won’t be able to properly record operations and data flows<br>
for that variable?</p>
</blockquote>
</aside>
<p>No, Autograd will still track all operations on parameters etc. However, the parameters (of the child modules) won’t be registered internally, so that <code>model.parameters()</code> won’t return them (they might be thus missing in when you pass the parameters with this call to the optimizer). Also <code>model.to()</code> won’t grab these parameters and they will stay on the initial device.</p>

hongtaesuk: <p>Woah…!<br>
goosebumps!! haha<br>
model.to() didn’t grab those parameters and it actually arose a operation confliction<br>
in two different devices (cpu &amp; gpu).</p>
<p>I don’t think I understood the whole things you’ve said,<br>
but I guess what I produce as a parent module which holds the normal python dict,<br>
produces - somehow, what I didn’t know in advance - child modules internally<br>
but those normal python dicts are not registered properly,<br>
is this correct?</p>
<p>Sorry to bother you.</p>

ptrblck: <p>Yes, your explanation is correct. <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=9" title=":slight_smile:" class="emoji" alt=":slight_smile:"><br>
Therefore you could still use lists and dicts, if you explicitly don’t want to register the content inside the module.<br>
A bit unrelated to this, but <code>self.param = nn.Parameter(...)</code> and <code>self.register_buffer</code> will register the parameters/buffers properly, while a simple assignment of a tensor (<code>self.tensor = torch.tensor(...)</code>) will not.</p>
<p>Feel free to ask, in case something is unclear or you want more information.</p>

hongtaesuk: <p>Awesome!!!<br>
so good to get such a quality answers <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=9" title=":slight_smile:" class="emoji" alt=":slight_smile:"><br>
Thanks a lot ptrblck!!<br>
and merry Christmas</p>