Unexpected outputs for forward hooks

vdw: <p>To get a bit more familiar with hooks (just forward hooks for now), I’ve created a very simple LSTM model:</p>
<pre><code class="lang-auto">class SimpleLSTM(nn.Module):
    
    def __init__(self, vocab_size, out_size):
        super().__init__()
        
        self.embed = nn.Embedding(vocab_size, 5)
        self.lstm = nn.LSTM(5, 4, bidirectional=True)
        self.linear1 = nn.Linear(4*2, 4)
        self.relu1 = nn.ReLU()
        self.dropout1 = nn.Dropout(0.5)
        self.linear2 = nn.Linear(4, out_size)
        self.log_softmax = nn.LogSoftmax(dim=1)

    def forward(self, X):
        batch_size = X.shape[0]
        out = self.embed(X)
        out = out.transpose(0, 1)
        out, (h, c) = self.lstm(out)
        h = h.view(1, 2, batch_size, 4)[-1]
        h = torch.concat((h[0], h[1]), dim=1)
        out = self.linear1(h)
        out = self.relu1(out)
        out = self.dropout1(out)
        out = self.linear2(out)
        out = self.log_softmax(out)
        return(out)
</code></pre>
<p>I’ve then attached forward hooks to each layer, with the sole purpose to print the input shapes of the output shapes. When I run a sample batch of <code>(3,5)</code> – meaning <code>(batch_size, seq_len)</code>, I get the following outputs:</p>
<pre><code class="lang-auto">------- Embedding(10, 5) -------
embed [ in] torch.Size([3, 5])
embed [out] torch.Size([5, 5])

------- LSTM(5, 4, bidirectional=True) -------
lstm [ in] torch.Size([5, 3, 5])
lstm [out] torch.Size([5, 3, 8])

------- Linear(in_features=8, out_features=4, bias=True) -------
linear1 [ in] torch.Size([3, 8])
linear1 [out] torch.Size([4])

------- ReLU() -------
relu1 [ in] torch.Size([3, 4])
relu1 [out] torch.Size([4])

------- Dropout(p=0.5, inplace=False) -------
dropout1 [ in] torch.Size([3, 4])
dropout1 [out] torch.Size([4])

------- Linear(in_features=4, out_features=2, bias=True) -------
linear2 [ in] torch.Size([3, 4])
linear2 [out] torch.Size([2])

------- LogSoftmax(dim=1) -------
log_softmax [ in] torch.Size([3, 2])
log_softmax [out] torch.Size([2])
</code></pre>
<p>It looks like the output shape is always missing the batch dimensions. The only exception is the <code>lstm</code> layer; I assume because <code>batch_first=False</code>.</p>
<p>Why is the output of a forward hook not the complete tensor? Or what might I do wrong here?</p>

ptrblck: <p>I’m not sure how your forward hook looks like, but I would guess you might be indexing into <code>input</code> and <code>output</code> in a way which might drop the batch dimension.<br>
I’m using a small and naive helper util. to flatten the potentially nested input/output tuples and get this output using your code:</p>
<pre><code class="lang-python">class SimpleLSTM(nn.Module):
    def __init__(self, vocab_size, out_size):
        super().__init__()
        self.embed = nn.Embedding(vocab_size, 5)
        self.lstm = nn.LSTM(5, 4, bidirectional=True)
        self.linear1 = nn.Linear(4*2, 4)
        self.relu1 = nn.ReLU()
        self.dropout1 = nn.Dropout(0.5)
        self.linear2 = nn.Linear(4, out_size)
        self.log_softmax = nn.LogSoftmax(dim=1)

    def forward(self, X):
        batch_size = X.shape[0]
        out = self.embed(X)
        out = out.transpose(0, 1)
        out, (h, c) = self.lstm(out)
        h = h.view(1, 2, batch_size, 4)[-1]
        h = torch.concat((h[0], h[1]), dim=1)
        out = self.linear1(h)
        out = self.relu1(out)
        out = self.dropout1(out)
        out = self.linear2(out)
        out = self.log_softmax(out)
        return(out)

def flatten(tensors):
    if isinstance(tensors, tuple):
        if len(tensors) == 0:
            return ()
        else:
            return flatten(tensors[0]) + flatten(tensors[1:])
    else:
        return (tensors,)

def get_hook(name):
    def hook(m, input, output):
        print(f"layer {name}")
        for i in flatten(input):
            print(f"input shape {i.shape}")
        for o in flatten(output):
            print(f"output shape {o.shape}")
    return hook
    

model = SimpleLSTM(vocab_size=10, out_size=20)
for name, module in model.named_modules():
    module.register_forward_hook(get_hook(name))

x = torch.randint(0, 10, (3, 5))
out = model(x)
# layer embed
# input shape torch.Size([3, 5])
# output shape torch.Size([3, 5, 5])
# layer lstm
# input shape torch.Size([5, 3, 5])
# output shape torch.Size([5, 3, 8])
# output shape torch.Size([2, 3, 4])
# output shape torch.Size([2, 3, 4])
# layer linear1
# input shape torch.Size([3, 8])
# output shape torch.Size([3, 4])
# layer relu1
# input shape torch.Size([3, 4])
# output shape torch.Size([3, 4])
# layer dropout1
# input shape torch.Size([3, 4])
# output shape torch.Size([3, 4])
# layer linear2
# input shape torch.Size([3, 4])
# output shape torch.Size([3, 20])
# layer log_softmax
# input shape torch.Size([3, 20])
# output shape torch.Size([3, 20])
# layer 
# input shape torch.Size([3, 5])
# output shape torch.Size([3, 20])
</code></pre>
<p>which should show all batch dimensions.</p>

vdw: <p>Right, both inputs and outputs come in different types. For the LSTM output it should have been obvious but I missed that as well.</p>
<p>It’s still feels a bit unintuitive that the inputs are always tuples, while the outputs are tensors for <code>nn.Linear</code>, <code>nn.Embedding</code>, etc. but tuples for <code>nn.LSTM</code> (again, this one makes sense of course):</p>
<pre><code class="lang-auto">------- Embedding(10, 5) -------
&lt;class 'tuple'&gt;
&lt;class 'torch.Tensor'&gt;

------- LSTM(5, 4, bidirectional=True) -------
&lt;class 'tuple'&gt;
&lt;class 'tuple'&gt;

------- Linear(in_features=8, out_features=4, bias=True) -------
&lt;class 'tuple'&gt;
&lt;class 'torch.Tensor'&gt;

------- ReLU() -------
&lt;class 'tuple'&gt;
&lt;class 'torch.Tensor'&gt;

------- Dropout(p=0.5, inplace=False) -------
&lt;class 'tuple'&gt;
&lt;class 'torch.Tensor'&gt;

------- Linear(in_features=4, out_features=2, bias=True) -------
&lt;class 'tuple'&gt;
&lt;class 'torch.Tensor'&gt;

------- LogSoftmax(dim=1) -------
&lt;class 'tuple'&gt;
&lt;class 'torch.Tensor'&gt;
</code></pre>
<p>But yes, since thought it’s always a tuple, I used</p>
<pre><code class="lang-auto">print(input[0])
print(output[0])
</code></pre>
<p>which is fine for the inputs (always a tuple), but only prints the first sample in the batch if the output is a tensor.</p>