batch_first=True gives different results in RNN

kinwai_cheuk: <p>I am trying to rewrite some code to make them batch_first=True.<br>
Then, I realized that the results when setting <code>batch_first=True</code> is different from that of <code>batch_first=False</code> even if I am using the same random seed.</p>
<p>Here’s my original code written with <code>batch_first=False</code></p>
<pre><code class="lang-auto">import torch
import torch.nn as nn
import torch.optim as optim
from torchvision.datasets import MNIST
from torchvision import transforms
from torch.utils.data import DataLoader

import numpy as np

import random
import math
import time

SEED = 1234

random.seed(SEED)
np.random.seed(SEED)
torch.manual_seed(SEED)
torch.cuda.manual_seed(SEED)
torch.backends.cudnn.deterministic = True

device = 'cuda:1'

transform = transforms.Compose(
    [transforms.ToTensor(),
     transforms.Normalize((0.5, ), (0.5, ))])
trainset = MNIST(root='./data', 
                                        train=True,
                                        download=True,
                                        transform=transform)
trainloader = torch.utils.data.DataLoader(trainset,
                                            batch_size=100,
                                            shuffle=True,
                                            num_workers=2)

testset = MNIST(root='./data', 
                                        train=False, 
                                        download=True, 
                                        transform=transform)
testloader = torch.utils.data.DataLoader(testset, 
                                            batch_size=100,
                                            shuffle=False, 
                                            num_workers=2)

class Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.emb = nn.Linear(28,14)        
        self.rnn = nn.GRU(14, 14, batch_first=False)
        self.fc = nn.Linear(14,10)
        self.dropout = nn.Dropout(p=0.5)
        
    def forward(self, x):
        x = self.dropout(self.emb(x))
        outputs, hidden = self.rnn(x)
        output = torch.sigmoid(self.fc(hidden.squeeze(0)))

        return output

model = Model()
model = model.to(device)
optimizer = optim.Adam(model.parameters())
criterion = nn.CrossEntropyLoss()

for e in range(5):
    train_total_loss = 0
    model.train()
    for i, (x,y) in enumerate(trainloader):
        x = x.to(device)
        y = y.to(device)
        x = x.squeeze(1).transpose(0,1)
        optimizer.zero_grad()
        y_pred = model(x)
        loss = criterion(y_pred, y)
        loss.backward()
        optimizer.step()
        train_total_loss += loss.item()
        print(f'{i}/{len(trainloader)}: loss = {loss.item()}', end='\r')
        

    test_total_loss = 0
    model.eval()
    for i, (x,y) in enumerate(testloader):
        x = x.to(device)
        y = y.to(device)      
        x = x.squeeze(1).transpose(0,1)
        y_pred = model(x)
        loss = criterion(y_pred, y)
        test_total_loss += loss.item()
        print(f'{i}/{len(testloader)}: loss = {loss.item()}', end='\r')        
    print(f'epoch {e+1}: train loss={train_total_loss/len(trainloader):.4f}\ttest loss={test_total_loss/len(testloader):.4f}')

output &gt;&gt;&gt;
epoch 1: train loss=2.1462	test loss=1.9462
epoch 2: train loss=1.9133	test loss=1.8628
epoch 3: train loss=1.8353	test loss=1.7833
epoch 4: train loss=1.7718	test loss=1.7382
epoch 5: train loss=1.7444	test loss=1.7191
</code></pre>
<p>Now I change my model to <code>batch_first=False</code></p>
<pre><code class="lang-auto">class Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.emb = nn.Linear(28,14)        
        self.rnn = nn.GRU(14, 14, batch_first=True)
        self.fc = nn.Linear(14,10)
        self.dropout = nn.Dropout(p=0.5)
        
    def forward(self, x):
        

        
        x = self.dropout(self.emb(x))
        outputs, hidden = self.rnn(x)
        output = torch.sigmoid(self.fc(hidden.squeeze(0)))
     
        return output

model = Model()
model = model.to(device)
optimizer = optim.Adam(model.parameters())
criterion = nn.CrossEntropyLoss()

for e in range(5):
    train_total_loss = 0
    model.train()
    for i, (x,y) in enumerate(trainloader):
        x = x.to(device)
        y = y.to(device)
        optimizer.zero_grad()
        y_pred = model(x.squeeze(1))
        loss = criterion(y_pred, y)
        loss.backward()
        optimizer.step()
        train_total_loss += loss.item()
        print(f'{i}/{len(trainloader)}: loss = {loss.item()}', end='\r')
        

    test_total_loss = 0
    model.eval()
    for i, (x,y) in enumerate(testloader):
        x = x.to(device)
        y = y.to(device)        
        y_pred = model(x.squeeze(1))
        loss = criterion(y_pred, y)
        test_total_loss += loss.item()
        print(f'{i}/{len(testloader)}: loss = {loss.item()}', end='\r')        
    print(f'epoch {e+1}: train loss={train_total_loss/len(trainloader):.4f}\ttest loss={test_total_loss/len(testloader):.4f}')

output&gt;&gt;&gt;
epoch 1: train loss=2.1360	test loss=1.9413
epoch 2: train loss=1.9016	test loss=1.8367
epoch 3: train loss=1.8216	test loss=1.7775
epoch 4: train loss=1.7840	test loss=1.7494
epoch 5: train loss=1.7612	test loss=1.7350
</code></pre>
<p>The result is different.<br>
What is wrong here?</p>

ptrblck: <p>I cannot reproduce the difference for the <code>nn.GRU</code> module in isolation:</p>
<pre><code class="lang-python">device = 'cuda'
x0 = torch.randn(5, 6, 14, device=device)
model0 = nn.GRU(14, 14, batch_first=True).to(device)
out0 = model0(x0)

x1 = x0.permute(1, 0, 2)
model1 = nn.GRU(14, 14, batch_first=False).to(device)
model1.load_state_dict(model0.state_dict())
out1 = model1(x1)

print((out0[0].permute(1, 0, 2) - out1[0]).abs().max())
&gt; tensor(0., device='cuda:0', grad_fn=&lt;MaxBackward1&gt;)
print((out0[1] - out1[1]).abs().max())
&gt; tensor(0., device='cuda:0', grad_fn=&lt;MaxBackward1&gt;)
</code></pre>
<p>So the difference is most likely coming from another part of your code.<br>
I would recommend to slim down the code by e.g. removing the data loading etc. to get rid of potential sources which could add randomness.<br>
Also, since your model is using dropout you have to be very careful to seed to code in the appropriate place and make sure the same calls to the pseudo-random number generator are performed in both scripts. Thus I would just remove it or call <code>model.eval()</code>.</p>

kinwai_cheuk: <p>Thanks for your insight. Indeed, the problem comes from the <code>dropout</code> layer.<br>
Even when I make sure the seed for both script are the same</p>
<pre><code class="lang-auto">import torch
import torch.nn as nn
import numpy as np
import random
SEED = 1234
random.seed(SEED)
np.random.seed(SEED)
torch.manual_seed(SEED)
torch.cuda.manual_seed(SEED)
torch.backends.cudnn.deterministic = True

x_batch_first = torch.ones(2,3,4)
x_batch_second = torch.ones(3,2,4)

dropout_layer = nn.Dropout(p=0.5)

output1 = dropout_layer(x_batch_first)
output2 = dropout_layer(x_batch_second)
torch.equal(output1, output2.transpose(0,1))
output &gt;&gt;&gt; False
</code></pre>
<p>When I print out both results, it gives me:</p>
<pre><code class="lang-auto">output1
tensor([[[0., 2., 0., 0.],
         [2., 0., 2., 0.],
         [0., 0., 2., 2.]],

        [[2., 0., 2., 2.],
         [0., 0., 2., 2.],
         [0., 2., 2., 0.]]])

output2
tensor([[[2., 0., 2., 0.],
         [2., 2., 0., 2.]],

        [[0., 2., 2., 2.],
         [2., 0., 2., 0.]],

        [[2., 2., 0., 2.],
         [2., 0., 0., 0.]]])
</code></pre>
<p>Any ways to make sure dropout layer still behaves the same when <code>batch_first=True</code> and <code>batch_first=False</code>?</p>

ptrblck: <p>The output of your current script is expected, since each call into <code>dropout_layer</code> will sample a new random mask.<br>
The initial seed makes sure that sequential calls into the random number generator will yield the same “random” numbers, but of course it’s not forcing dropout layers to sample a constant mask.</p>
<p>If you want to force the dropout layers to behave exactly the same, you could either seed the code right before the call (which would get rid of the randomness and thus I would assume your model to perform badly) via:</p>
<pre><code class="lang-python">x_batch_first = torch.ones(2,3,4)
x_batch_second = torch.ones(3,2,4)

dropout_layer = nn.Dropout(p=0.5)

torch.manual_seed(2809)
output1 = dropout_layer(x_batch_first)
torch.manual_seed(2809)
output2 = dropout_layer(x_batch_second.transpose(0, 1))
torch.equal(output1, output2)
&gt; True
</code></pre>
<p>or maybe sample the dropout masks manually outside of the training script and just apply them.</p>
<p>Note that even seeding the layer would give different results, if the tensors have different shapes, which is why I’ve transposed the second input before the dropout call.</p>