PyTorch model ported from Keras model not learning

shivin7: <p>I am porting a net from Keras to Pytorch, however the training in Pytorch doesn’t seem to learn anything. The net is trying to learn sub-word embeddings in a phrase and classify it among three classes.</p>
<p>The architecture (along with output shapes) of the net is as follows :</p>
<ol>
<li>Embedding Layer (batch x 200 x 128)</li>
<li>Convolution Layer (batch x 198 x 128)</li>
<li>Max Pooling Layer (batch x 66 x 128)</li>
<li>LSTM Layer (batch x 66 x 128)</li>
<li>LSTM Layer (batch x 128)</li>
<li>Dense Layer (batch x 3)</li>
</ol>
<p>I have tried with different learning rates, changing the loss from Cross Entropy to Softmax+NLL, however the Pytorch model doesn’t seem to learn.</p>
<p>The input <code>X</code> is a file of shape <code>(num_samples, 200)</code> with each value in the range [0,26], and <code>y</code> are the labels in [0,1,2] for each sample.</p>
<p>Here is the Keras code :</p>
<pre><code class="lang-auto">from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Dense, Dropout, Activation, Embedding, LSTM, GRU, Convolution1D, MaxPooling1D
from tensorflow.keras.utils import to_categorical

#Format conversion for y
y_train = to_categorical(y, 3) 
#Train &amp; Validation data splitting
X_train, X_valid, y_train, y_valid = train_test_split(X, y, test_size=0.2, random_state=42)

embedLayer = Embedding(input_dim=27, output_dim=128, input_length=200)
convLayer = Convolution1D(filters=128, kernel_size=3, activation='relu')
poolLayer = MaxPooling1D(pool_size=3)
lstmLayer1 = LSTM(units=128, dropout=0.2, recurrent_dropout=0.2, return_sequences=True)
lstmLayer2 = LSTM(units=128, dropout=0.2, recurrent_dropout=0.2, return_sequences=False)
denseLayer = Dense(units=3)

model = Sequential()
model.add(embedLayer)
model.add(convLayer)
model.add(poolLayer)
model.add(lstmLayer1)
model.add(lstmLayer2)
model.add(denseLayer)
model.add(Activation('softmax'))

model.compile(loss='categorical_crossentropy',
              optimizer='adamax',
              metrics=['accuracy'])

# Training data
model.fit(x=X_train, y=y_train, 
          batch_size=128, 
          epochs=50,
          validation_data=(X_valid, y_valid))
</code></pre>
<p>Here is the PyTorch code :</p>
<pre><code class="lang-auto">class RNNModel(nn.Module):
    
    def __init__(self):
        
        super(RNNModel, self).__init__()
            
        # Layers
        self.embeddingLayer = nn.Embedding(num_embeddings=27, embedding_dim=128)
        self.convLayer = nn.Conv1d(in_channels=128, out_channels=128, 
                                   kernel_size=3)
        self.lstmLayer1 = nn.LSTM(input_size=128, hidden_size=128)
        self.lstmLayer2 = nn.LSTM(input_size=128, hidden_size=128)
        self.denseLayer = nn.Linear(in_features=128, out_features=3)
        
    def forward(self, x):
        
        # Forming embeddings
        x = self.embeddingLayer(x)
        
        # Convolution and pooling
        x = x.view(-1, 128, 200)
        x = F.relu(self.convLayer(x))
        x = F.max_pool1d(x, kernel_size=3)
        
        # LSTM layers
        x = x.view(x.shape[2], x.shape[0], x.shape[1])
        x, _ = self.lstmLayer1(x)
        x = F.dropout(x, p=0.2)
        _, (x, _) = self.lstmLayer2(x)
        x = F.dropout(x, p=0.2)
        
        # Dense layer
        x = x.view(x.shape[1], x.shape[2])
        x = self.denseLayer(x)
        
        return x
        
device = torch.device('cuda')

model = RNNModel()
model = model.to(device)

lossFunc = nn.CrossEntropyLoss()
optimizer = torch.optim.Adamax(model.parameters())

# trainGen and devGen are generators for training and dev set respectively

for i in range(50):
    for X_batch, y_batch in trainGen:
        X_batch = X_batch.to(device)
        y_batch = y_batch.to(device)
        pred_batch = model(X_batch.long())
        loss = lossFunc(pred_batch, y_batch.long())
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
    print("\nITERATION ", i+1, "\n Train Loss = ", loss.item())
    
    with torch.no_grad():
        for X_dev, y_dev in devGen:
            X_dev = X_dev.to(device)
            y_dev = y_dev.to(device)
            pred_batch_dev = model(X_dev.long())
            loss = lossFunc(pred_batch_dev, y_dev.long())
            print(" Dev Loss = ", loss.item())
</code></pre>
<p>The Keras model does learn something because it brings up the accuracy from around 47% to 63%. However the Pytorch models keeps wavering around the initial loss only. What is the problem here?</p>

ptrblck: <p>Could you explain, what <code>dropout=0.2</code> and <code>recurrent_dropout=0.2</code> do in your Keras model?<br>
It seems that you are using a single <code>F.dropout</code> on the outputs of both LSTM layers.</p>
<p>Also, since you are using the functional dropout API, you should pass the <code>training=self.training</code> argument to disable it during validation. <img src="https://discuss.pytorch.org/images/emoji/apple/wink.png?v=9" title=":wink:" class="emoji" alt=":wink:"></p>

shivin7: <aside class="quote no-group">
<blockquote>
<p>Could you explain, what <code>dropout=0.2</code> and <code>recurrent_dropout=0.2</code> do in your Keras model?<br>
It seems that you are using a single <code>F.dropout</code> on the outputs of both LSTM layers.</p>
</blockquote>
</aside>
<p>From the keras docs, <code>dropout</code> is for the linear transformation of the inputs and <code>recurrent_dropout</code> is for the linear transformation of the recurrent states.<br>
I’m guessing keras <code>dropout</code> masks the inputs, hence that should simply be implemented by adding a <code>F.dropout</code> layer before each LSTM in PyTorch (I’ve mistakenly added them at the outputs, I’ll change it right away)<br>
I’m not able to understand how to implement the <code>recurrent_dropout</code> part in PyTorch, which masks the connections between the recurrent units. Is it possible that absence of this particular dropout is the only reason why the PyTorch model is not learning while Keras one is?</p>
<aside class="quote no-group">
<blockquote>
<p>Also, since you are using the functional dropout API, you should pass the <code>training=self.training</code> argument to disable it during validation. <img src="https://discuss.pytorch.org/images/emoji/apple/wink.png?v=12" title=":wink:" class="emoji" alt=":wink:" loading="lazy" width="20" height="20"></p>
</blockquote>
</aside>
<p>Oh yes, my bad. I’ll change it right away.</p>
<p>Thanks a lot!</p>

shivin7: <p>So I read up a little and it seems that recurrent dropout is not yet implemented in PyTorch, and there is an open <a href="https://github.com/pytorch/pytorch/issues/9572" rel="nofollow noopener">Github issue</a> for the same. I guess the only solution is to manually write the LSTM cell then.<br>
However, shouldn’t absence of such a dropout be responsible for overfitting at max?<br>
I still can’t wrap my head around the fact that this is the sole reason why the model isn’t learning at all (neither the training nor the dev loss seem to decrease at all).</p>

ptrblck: <aside class="quote no-group" data-username="shivin7" data-post="4" data-topic="77192">
<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/shivin7/48/22921_2.png" class="avatar"> shivin7:</div>
<blockquote>
<p>I guess the only solution is to manually write the LSTM cell then. However, shouldn’t absence of such a dropout be responsible for overfitting at max?</p>
</blockquote>
</aside>
<p>That might be the case and I agree that something else might create the issue.</p>
<p>In the Keras implementation you are using <code>return_sequences=False</code>.<br>
Based on the docs, it seems that this LSTM will only return the output of the last timestep.</p>
<p>In your PyTorch implementation you are reshaping the output of <code>self.lstmLayer2</code> via:</p>
<pre data-code-wrap="python"><code class="lang-python">x = x.view(x.shape[1], x.shape[2])
</code></pre>
<p>which is not using the last timestep, but the complete sequence.</p>
<p>That being said, I think the reshaping in general might be wrong.</p>
<p>The last pooling layer will return an activation in the shape <code>[batch_size, channels, seq_len]</code>.<br>
Since the <code>nn.LSTM</code> will accept an input of <code>[seq_len, batch_size, nb_features]</code> by default (if <code>batch_first=False</code>).<br>
It seems you would like to permute the activation here:</p>
<pre data-code-wrap="python"><code class="lang-python">x = x.view(x.shape[2], x.shape[0], x.shape[1])
</code></pre>
<p>which won’t work, since you are interleaving the signal.<br>
Use <code>x = x.permute(2, 0, 1)</code> instead to permute the dimensions.</p>
<p>The same applies for the output of <code>self.lstmLayer2</code>.</p>
<p>Let me know, if that helps.</p>

shivin7: <aside class="quote no-group" data-username="ptrblck" data-post="5" data-topic="77192">
<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/ptrblck/48/1823_2.png" class="avatar"> ptrblck:</div>
<blockquote>
<p>which is not using the last timestep, but the complete sequence.</p>
</blockquote>
</aside>
<p>But I’m taking the output of the second LSTM as <code>_, (x, _) = self.lstmLayer2(x)</code>, which according to the docs produces <code>output, (h_n, c_n)</code> that is the output of the last timestep.<br>
The  <code>x = x.view(x.shape[1], x.shape[2])</code> after the second LSTM was to reshape <code>(seq_len, batch)</code> into <code>(batch, seq_len)</code> to be fed into the dense layer.</p>
<aside class="quote no-group">
<blockquote>
<p>Use <code>x = x.permute(2, 0, 1)</code> instead to permute the dimensions.</p>
</blockquote>
</aside>
<p>Yes! I changed the <code>view</code> reshaping before <code>convLayer</code> and <code>lstmLayer1</code> to <code>permute</code> and that got it working. Thanks a lot!<br>
This works perfectly now :</p>
<pre><code class="lang-auto">    def forward(self, x):
        
        # Forming embeddings
        x = self.embeddingLayer(x)
        
        # Convolution and pooling
        x = x.permute(0, 2, 1)
        x = F.relu(self.convLayer(x))
        x = F.max_pool1d(x, kernel_size=pool_length)
        
        # LSTM layers
        x = x.permute(2, 0, 1)
        x = F.dropout(x, p=0.2, training=self.training)
        x, _ = self.lstmLayer1(x)
        x = F.dropout(x, p=0.2, training=self.training)
        _, (x, _) = self.lstmLayer2(x)
        
        # Dense layer
        x = x.view(x.shape[1], x.shape[2])
        x = self.denseLayer(x)
        
        return x
</code></pre>
<p>But I can’t understand why <code>view</code> doesn’t work here. What do you mean by interleaving the signal?</p>

ptrblck: <p>Thanks for the correction regarding the last state / output, as I’ve clearly mixed up the return values. <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>
<aside class="quote no-group" data-username="shivin7" data-post="6" data-topic="77192">
<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/shivin7/48/22921_2.png" class="avatar"> shivin7:</div>
<blockquote>
<p>But I can’t understand why <code>view</code> doesn’t work here. What do you mean by interleaving the signal?</p>
</blockquote>
</aside>
<p>Have a look at this example:</p>
<pre data-code-wrap="python"><code class="lang-python"> x = torch.arange(2*4*5).view(2, 4, 5)
print(x)

# Now permute the dimensions
print(x.permute(0, 2, 1)) # rows and columns for each sample in dim0 are swapped

# Wrong permutation!
print(x.view(2, 5, 4)) # result is interleaved
</code></pre>
<p>Here you can see, that I would like to swap the last two dimensions (basically transpose dim1 and dim2 for each sample in the batch dimension).<br>
<code>.permute</code> does the job right, while <code>view</code> just changes the stride and shape and thus will create a wrong and interleaved result.<br>
By “interleaved” I mean that the original row tensors are cut and the remaining elements are added to the next row.</p>

shivin7: <p>Oh alright, got the point!<br>
Thanks a lot for the help! <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=9" title=":slight_smile:" class="emoji" alt=":slight_smile:"></p>