ModuleDict to route observations to different heads

username: ilkarman
cooked: <p>I would like to route different observations within a batch to a different pre-processing head (small NN), then all of them go through the same backbone-CNN. I thought I could achieve this like so using ModuleDict:</p>
<pre><code class="lang-auto">class NNSandwich(nn.Module):
    def __init__(self, num_heads):
        super(NNSandwich, self).__init__()
        
        # Pre-NN
        self.pre_choices = nn.ModuleDict()
        for i in range(num_heads):
                self.pre_choices["pre_{}".format(i)] = nn.Sequential(
                    nn.Conv2d(3, 32, 3),
                    nn.ReLU(),
                    nn.Conv2d(32, 3, 3),
                    nn.ReLU()
                )    
                
        # Shared-backbone
        self.shared_backbone = torchvision.models.resnet18(True)
        self.fc = nn.Linear(1000, 2)
        
    def forward(self, x, choices):
         
        x = torch.cat(
            [self.pre_choices[j](x[i].unsqueeze(0)) for i,j in enumerate(choices)],
            dim=0
        )
        x = self.shared_backbone(x)
        x = self.fc(x)
        return x
</code></pre>
<p>However, if seems that choice can’t be a list (different for each observation) but a string that is thus the same for each observation. Thus I use list comprehension and torch.cat() to pass through the respective heads.</p>
<p>However, this method doesn’t seem to update the weights properly:</p>
<pre><code class="lang-auto">import torch
import torch.nn as nn
import torchvision
import random

HEADS = 5
BATCH = 4

# NNSandwich
sandwich_nn = NNSandwich(num_heads=HEADS)
sandwich_nn.cuda()

# Optimiser
optimizer = torch.optim.SGD(sandwich_nn.parameters(), lr=0.5)

# Input
batch_i = torch.ones(BATCH,3,224,224).cuda()
labels_i = torch.ones(BATCH).cuda().long()
choices_i = ['pre_2', 'pre_1', 'pre_0', 'pre_4']

print(sandwich_nn.pre_choices.pre_0[0].weight[0][0])
print(sandwich_nn.pre_choices.pre_1[0].weight[0][0])
print(sandwich_nn.pre_choices.pre_2[0].weight[0][0])
print(sandwich_nn.pre_choices.pre_3[0].weight[0][0])
print(sandwich_nn.pre_choices.pre_4[0].weight[0][0])

# Forwards-Pass
out = sandwich_nn(batch_i, choices_i)
print(out)
# Backwards-Pass
loss = nn.functional.cross_entropy(out, labels_i)
optimizer.zero_grad()
loss.backward()
optimizer.step()

print(sandwich_nn.pre_choices.pre_0[0].weight[0][0])
print(sandwich_nn.pre_choices.pre_1[0].weight[0][0])
print(sandwich_nn.pre_choices.pre_2[0].weight[0][0])
print(sandwich_nn.pre_choices.pre_3[0].weight[0][0])
print(sandwich_nn.pre_choices.pre_4[0].weight[0][0])
</code></pre>
<p>In this case only <code>print(sandwich_nn.pre_choices.pre_1[0].weight[0][0])</code> differs, it seems the other heads haven’t been updated (but pre2, pre1, pre0, pre4 should be)</p>

username: ptrblck
cooked: <p>Your code should work and it seems the gradients are just “killed” by the <code>ReLU</code>s in <code>pre_choices</code>.<br>
I.e. if you remove both <code>nn.ReLU()</code> non-linearities, you’ll see valid updated, as they might create zero gradients for some iterations.<br>
This doesn’t necessarily mean that your model won’t get any valid gradients in the following iterations.</p>
<p>Also, if you want to use an integer index instead of <code>keys</code>, you could use <code>nn.ModuleList</code> instead of <code>nn.ModuleDict</code>.</p>

username: ilkarman
cooked: <p>Thank you for the reply Patrick!</p>
<p>I’m still a bit confused to understand however how my gradient will be calculated w.r.t the observations in the batch with the for-loop.</p>
<p>For example if my nn.Module has ModuleDict(Block1A, Block1B, block1C) -&gt; Block2 -&gt; Block3 -&gt; Out</p>
<p>And in my input-batch I have say (all) observations corresponding to Block1A then wouldn’t the following produce gradients:</p>
<p>One-by-one:<br>
<code>x = torch.cat([self.moduledict['Block1A'](x[i].unsqueeze(0)) for i in range(len(x))])</code></p>
<p>Together:<br>
<code>x = self.moduledict['Block1A'](x)</code></p>
<p>Just to illustrate with a code-example:</p>
<pre><code class="lang-auto">import torch
torch.manual_seed(0)
import torch.nn as nn
import torch.nn.functional as F

class Flatten(nn.Module):
    def forward(self, input):
        return input.view(input.size(0), -1)

class Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.pre_heads = nn.ModuleDict()
        for head_name in ['a','b','c']:
                self.pre_heads[head_name] = nn.Sequential(
                    nn.Conv2d(16,16, kernel_size=1),
                    nn.ReLU()
                )
        self.out = nn.Sequential(
            Flatten(),
            nn.Linear(16*4*4,1)
        )
        
    def forward(self, x, head_name):           
        x = torch.cat(
            [self.pre_heads[j](x[i].unsqueeze(0)) for i, j in enumerate(head_name)],
            dim=0
        )
        x = torch.squeeze(torch.sigmoid(self.out(x)))
        return x

mymodel = Model()
sample_input = torch.ones(4,16,4,4)
output = mymodel(sample_input, head_name=['a', 'a', 'a', 'a'])
loss = F.binary_cross_entropy(output, torch.Tensor([1,1,1,1]))
loss.backward()
mymodel.pre_heads['a'][0].weight.grad.sum()
# tensor(2.3628)
</code></pre>
<p>v.s.</p>
<pre><code class="lang-auto">class Model(nn.Module):        
    # Only changed Forward
    def forward(self, x):           
        x = self.pre_heads['a'](x)
        x = torch.squeeze(torch.sigmoid(self.out(x)))
        return x

mymodel = Model()
sample_input = torch.ones(4,16,4,4)
output = mymodel(sample_input])
loss = F.binary_cross_entropy(output, torch.Tensor([1,1,1,1]))
loss.backward()
mymodel.pre_heads['a'][0].weight.grad.sum()
# tensor(-0.9457)
</code></pre>

username: ptrblck
cooked: <p>Both approaches should yield the same result and I assume your models might have used different layer parameters.<br>
This code snippet shows that the different <code>forward</code> implementations create the same gradients up the floating point precision:</p>
<pre><code class="lang-python">import torch
torch.manual_seed(0)
import torch.nn as nn
import torch.nn.functional as F

class Flatten(nn.Module):
    def forward(self, input):
        return input.view(input.size(0), -1)

class Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.pre_heads = nn.ModuleDict()
        for head_name in ['a','b','c']:
                self.pre_heads[head_name] = nn.Sequential(
                    nn.Conv2d(16,16, kernel_size=1),
                    nn.ReLU()
                )
        self.out = nn.Sequential(
            Flatten(),
            nn.Linear(16*4*4,1)
        )
    
    def forward(self, x, head_name, use_heads=True):
        if use_heads:
            x = torch.cat(
                [self.pre_heads[j](x[i].unsqueeze(0)) for i, j in enumerate(head_name)],
                dim=0
            )
            x = torch.squeeze(torch.sigmoid(self.out(x)))
            return x
        else:
            x = self.pre_heads['a'](x)
            x = torch.squeeze(torch.sigmoid(self.out(x)))
            return x

mymodel = Model()
sample_input = torch.ones(4,16,4,4)

# 
output = mymodel(sample_input, head_name=['a', 'a', 'a', 'a'])
loss = F.binary_cross_entropy(output, torch.Tensor([1,1,1,1]))
loss.backward()
grad0 = mymodel.pre_heads['a'][0].weight.grad.sum().clone()

#
mymodel.zero_grad()
output = mymodel(sample_input, head_name=None, use_heads=False)
loss = F.binary_cross_entropy(output, torch.Tensor([1,1,1,1]))
loss.backward()
grad1 = mymodel.pre_heads['a'][0].weight.grad.sum().clone()

print((grad0 - grad1).abs())
&gt; tensor(1.1921e-07)
</code></pre>

username: ilkarman
cooked: <p>Thank you for testing this out! It makes sense to me now (I guess I did have the two models differently initialised despite setting seed)</p>