400% higher error with PyTorch compared with identical Keras model (with Adam optimizer)

username: ProteinGuy
cooked: <hr>
<p><strong>TLDR</strong>:</p>
<p><em>A simple (single hidden-layer) feed-forward Pytorch model trained to predict the function <code>y = sin(X1) + sin(X2) + ... sin(X10)</code> substantially underperforms an identical model built/trained with Keras. Why is this so and what can be done to mitigate the difference in performance?</em></p>
<hr>
<p>In training a regression model, I noticed that PyTorch drastically underperforms an identical model built with Keras.</p>
<p><strong>This phenomenon has been observed and reported previously</strong>:</p>
<ul>
<li>
<p><a href="https://discuss.pytorch.org/t/the-same-model-produces-worse-results-on-pytorch-than-on-tensorflow/5380">The same model produces worse results on pytorch than on tensorflow</a></p>
</li>
<li>
<p><a href="https://discuss.pytorch.org/t/cnn-model-in-pytorch-giving-30-less-accuracy-to-tensoflowflow-model/85410">CNN model in pytorch giving 30% less accuracy to Tensoflowflow model</a>:</p>
</li>
<li>
<p><a href="https://discuss.pytorch.org/t/pytorch-adam-vs-tensorflow-adam/74471">PyTorch Adam vs Tensorflow Adam</a></p>
</li>
<li>
<p><a href="https://discuss.pytorch.org/t/suboptimal-convergence-when-compared-with-tensorflow-model/5099">Suboptimal convergence when compared with TensorFlow model</a></p>
</li>
<li>
<p><a href="https://discuss.pytorch.org/t/rnn-and-adam-slower-convergence-than-keras/11278">RNN and Adam: slower convergence than Keras</a></p>
</li>
<li>
<p><a href="https://discuss.pytorch.org/t/pytorch-comparable-but-worse-than-keras-on-a-simple-feed-forward-network/9928">PyTorch comparable but worse than keras on a simple feed forward network</a></p>
</li>
<li>
<p><a href="https://www.reddit.com/r/pytorch/comments/ox0g4e/why_is_the_pytorch_model_doing_worse_than_the/" rel="noopener nofollow ugc">Why is the PyTorch model doing worse than the same model in Keras even with the same weight initialization?</a></p>
</li>
<li>
<p><a href="https://stackoverflow.com/questions/59344571/why-keras-behave-better-than-pytorch-under-the-same-network-configuration" rel="noopener nofollow ugc">Why Keras behave better than Pytorch under the same network configuration?</a></p>
</li>
</ul>
<p><strong>The following explanations and suggestions have been made previously as well</strong>:</p>
<ol>
<li>
<p>Using the same decimal precision (32 vs 64): <a href="https://discuss.pytorch.org/t/the-same-model-produces-worse-results-on-pytorch-than-on-tensorflow/5380">1</a>, <a href="https://www.reddit.com/r/MachineLearning/comments/7nw67c/d_pytorch_are_adam_and_rmsprop_okay/" rel="noopener nofollow ugc">2</a>,</p>
</li>
<li>
<p>Using a CPU instead of a GPU: <a href="https://discuss.pytorch.org/t/the-same-model-produces-worse-results-on-pytorch-than-on-tensorflow/5380">1</a>,<a href="https://discuss.pytorch.org/t/rnn-and-adam-slower-convergence-than-keras/11278/2">2</a></p>
</li>
<li>
<p>Change <code>retain_graph=True</code> to <code>create_graph=True</code> in computing the 2nd derivative with <code>autograd.grad</code>: <a href="https://discuss.pytorch.org/t/pytorch-adam-vs-tensorflow-adam/74471">1</a></p>
</li>
<li>
<p>Check if keras is using a regularizer, constraint, bias, or loss function in a different way from pytorch: <a href="https://discuss.pytorch.org/t/suboptimal-convergence-when-compared-with-tensorflow-model/5099/2">1</a>,<a href="https://discuss.pytorch.org/t/pytorch-comparable-but-worse-than-keras-on-a-simple-feed-forward-network/9928/4">2</a></p>
</li>
<li>
<p>Ensure you are computing the validation loss in the same way: <a href="https://discuss.pytorch.org/t/suboptimal-convergence-when-compared-with-tensorflow-model/5099/3">1</a></p>
</li>
<li>
<p>Use the same initialization routine: <a href="https://discuss.pytorch.org/t/suboptimal-convergence-when-compared-with-tensorflow-model/5099/3">1</a>,<a href="https://stackoverflow.com/questions/59344571/why-keras-behave-better-than-pytorch-under-the-same-network-configuration" rel="noopener nofollow ugc">2</a></p>
</li>
<li>
<p>Training the pytorch model for longer epochs: <a href="https://discuss.pytorch.org/t/rnn-and-adam-slower-convergence-than-keras/11278">1</a></p>
</li>
<li>
<p>Trying several random seeds: <a href="https://discuss.pytorch.org/t/rnn-and-adam-slower-convergence-than-keras/11278/8">1</a></p>
</li>
<li>
<p>Ensure that <code>model.eval()</code> is called in validation step when training pytorch model: <a href="https://discuss.pytorch.org/t/pytorch-comparable-but-worse-than-keras-on-a-simple-feed-forward-network/9928">1</a></p>
</li>
<li>
<p>The main issue is with the Adam optimizer, not the initialization: <a href="https://www.reddit.com/r/MachineLearning/comments/7nw67c/d_pytorch_are_adam_and_rmsprop_okay/" rel="noopener nofollow ugc">1</a></p>
</li>
</ol>
<p>To understand this issue, I trained a simple two-layer neural network (much simpler than my original model) in Keras and PyTorch, using the same hyperparameters and initialization routines, and following all the recommendations listed above. However, the PyTorch model results in a mean squared error (MSE) that is  400% higher than the MSE of the Keras model.</p>
<p><strong>Here is my code:</strong></p>
<p><strong>1. Generate a reproducible dataset</strong></p>
<pre data-code-wrap="python"><code class="lang-python">
def get_data():

    np.random.seed(0)
    Xtrain = np.random.normal(0, 1, size=(7000,10))
    Xval = np.random.normal(0, 1, size=(700,10))
    ytrain = np.sum(np.sin(Xtrain), axis=-1)
    yval = np.sum(np.sin(Xval), axis=-1)
    scaler = MinMaxScaler()
    ytrain = scaler.fit_transform(ytrain.reshape(-1,1)).reshape(-1)
    yval = scaler.transform(yval.reshape(-1,1)).reshape(-1) 

    return Xtrain, Xval, ytrain, yval



class XYData(Dataset):
    
    def __init__(self, X, y):
        
        super(XYData, self).__init__()
        self.X = torch.tensor(X, dtype=torch.float32)
        self.y = torch.tensor(y, dtype=torch.float32)
        self.len = len(y)
         
    def __getitem__(self, index):
        
        return (self.X[index], self.y[index])


    def __len__(self):

        return self.len

# Data, dataset, and dataloader
Xtrain, Xval, ytrain, yval = get_data()
traindata = XYData(Xtrain, ytrain)
valdata = XYData(Xval, yval)
trainloader = DataLoader(dataset=traindata, shuffle=True, batch_size=32, drop_last=False)
valloader = DataLoader(dataset=valdata, shuffle=True, batch_size=32, drop_last=False)
</code></pre>
<p><strong>2. Build Keras and PyTorch models with identical hyperparameters and initialization methods</strong></p>
<pre data-code-wrap="python"><code class="lang-python">class TorchLinearModel(nn.Module):
    
    def __init__(self, input_dim=10, random_seed=0):
        
        super(TorchLinearModel, self).__init__()
        _ = torch.manual_seed(random_seed)
        self.hidden_layer = nn.Linear(input_dim,100)
        self.initialize_layer(self.hidden_layer)        
        self.output_layer = nn.Linear(100, 1)
        self.initialize_layer(self.output_layer)

    def initialize_layer(self, layer):
        
        _ = torch.nn.init.xavier_normal_(layer.weight)
        #_ = torch.nn.init.xavier_uniform_(layer.weight)
        _ = torch.nn.init.constant(layer.bias,0)
        
    def forward(self, x):
        x = self.hidden_layer(x)
        x = self.output_layer(x)
        return x




def mean_squared_error(ytrue, ypred):
    
    return torch.mean(((ytrue - ypred) ** 2))


def build_torch_model():

    torch_model = TorchLinearModel()
    optimizer = optim.Adam(torch_model.parameters(), 
                           betas=(0.9,0.9999),
                           eps=1e-7,
                           lr=1e-3,
                           weight_decay=0)
    return torch_model, optimizer




def build_keras_model():
    
    x = layers.Input(shape=10)
    z = layers.Dense(units=100, activation=None, use_bias=True, kernel_regularizer=None, 
                     bias_regularizer=None)(x)
    y = layers.Dense(units=1, activation=None, use_bias=True, kernel_regularizer=None, 
                     bias_regularizer=None)(z)
    keras_model = Model(x, y, name='linear')
    optimizer = Adam(learning_rate=1e-3, beta_1=0.9, beta_2=0.9999, epsilon=1e-7, 
                     amsgrad=False)
    
    keras_model.compile(optimizer=optimizer, loss='mean_squared_error')
    
    return keras_model




# Instantiate models
torch_model, optimizer = build_torch_model()
keras_model = build_keras_model()

</code></pre>
<p><strong>3. Train PyTorch model for 100 epochs:</strong></p>
<pre data-code-wrap="python"><code class="lang-python">
torch_trainlosses, torch_vallosses = [], []

for epoch in range(100):

    # Training
    losses = []
    _ = torch_model.train()
    
    for i, (x,y) in enumerate(trainloader):
        optimizer.zero_grad()                          
        ypred = torch_model(x)
        loss = mean_squared_error(y, ypred) 
        _ = loss.backward()
        _ = optimizer.step()
        losses.append(loss.item())
    torch_trainlosses.append(np.mean(losses))
    
    # Validation
    losses = []
    _ = torch_model.eval()

    with torch.no_grad():
        for i, (x, y) in enumerate(valloader):
            ypred = torch_model(x)
            loss = mean_squared_error(y, ypred) 
            losses.append(loss.item())
    torch_vallosses.append(np.mean(losses))
    
    print(f"epoch={epoch+1}, train_loss={torch_trainlosses[-1]:.4f}, val_loss={torch_vallosses[-1]:.4f}")
    
</code></pre>
<p><strong>4. Train Keras model for 100 epochs:</strong></p>
<pre data-code-wrap="python"><code class="lang-python">history = keras_model.fit(Xtrain, ytrain, sample_weight=None, batch_size=32, epochs=100, 
                    validation_data=(Xval, yval))
</code></pre>
<p><strong>5. Loss in training history</strong></p>
<pre data-code-wrap="python"><code class="lang-python">plt.plot(torch_trainlosses, color='blue', label='PyTorch Train')    
plt.plot(torch_vallosses, color='blue', linestyle='--', label='PyTorch Val')  
plt.plot(history.history['loss'], color='brown', label='Keras Train')
plt.plot(history.history['val_loss'], color='brown', linestyle='--', label='Keras Val')
plt.legend()
</code></pre>
<p><img src="https://discuss.pytorch.org/uploads/default/original/3X/3/0/30ae605eac3f260692eba410f5422115c49a304a.png" alt="image" data-base62-sha1="6WEw8H6DTXOFzcbYQgdke2OG7BE" width="378" height="250"></p>
<p><em>Keras records a much lower error in the training. Since this may be due to a difference in how Keras computes the loss, I calculated the prediction error on the validation set with sklearn.metrics.mean_squared_error</em></p>
<p><strong>6. Validation error after training</strong></p>
<pre data-code-wrap="python"><code class="lang-python">ypred_keras = keras_model.predict(Xval).reshape(-1)
ypred_torch = torch_model(torch.tensor(Xval, dtype=torch.float32))
ypred_torch = ypred_torch.detach().numpy().reshape(-1)


mse_keras = metrics.mean_squared_error(yval, ypred_keras)
mse_torch = metrics.mean_squared_error(yval, ypred_torch)
print('Percent error difference:', (mse_torch / mse_keras - 1) * 100)  # 438%

r_keras = pearsonr(yval, ypred_keras)[0] # 0.911
r_pytorch = pearsonr(yval, ypred_torch)[0]  # 0.577
print("r_keras:", r_keras)
print("r_pytorch:", r_pytorch)

plt.scatter(ypred_keras, yval); plt.title('Keras'); plt.show(); plt.close()
plt.scatter(ypred_torch, yval); plt.title('Pytorch'); plt.show(); plt.close()
</code></pre>
<pre data-code-wrap="python"><code class="lang-python">Percent error difference: 479.1312469426776
r_keras: 0.9115184443702814
r_pytorch: 0.21728812737220082
</code></pre>
<p><img src="https://discuss.pytorch.org/uploads/default/original/3X/e/d/ed7bf2808b002ba4af97574a72081ba3717f12fc.png" alt="image" data-base62-sha1="xSSIptMPZVaWf2k57ehXxBJb6Gw" width="372" height="264"></p>
<p><img src="https://discuss.pytorch.org/uploads/default/original/3X/b/a/ba5c888db7a695a702bce5d7d4aa8d55033f613b.png" alt="image" data-base62-sha1="qAD3O7ns0vjhNl2fuDiptvO99mb" width="373" height="264"></p>
<p><em>The correlation of predicted values with ground truth is <ins>0.912</ins> for Keras but <ins>0.217</ins> for Pytorch, and the error for Pytorch is <ins>479%</ins> higher!</em></p>
<p><strong>7. Other trials</strong><br>
I also tried:</p>
<ul>
<li>Lowering the learning rate for Pytorch (lr=1e-4), <strong>R increases from 0.217 to 0.576</strong>, but it’s still much worse than Keras (r=0.912).</li>
<li>Increasing the learning rate for Pytorch (lr=1e-2), <strong>R is worse at 0.095</strong></li>
<li>Training numerous times with different random seeds. The <strong>performance is roughly the same</strong>, regardless.</li>
<li>Trained for longer than 100 epochs. No improvement was observed!</li>
<li>Used <code>torch.nn.init.xavier_uniform_</code> instead of <code>torch.nn.init.xavier_normal_</code> in the initialization of the weights. R <strong>improves from 0.217 to 0.639</strong>, but it’s still worse than Keras (0.912).</li>
</ul>
<hr>
<p><strong>What can be done to ensure that the PyTorch model converges to a reasonable error comparable with the Keras model?</strong></p>
<hr>

username: ptrblck
cooked: <p>You are broadcasting the prediction and target tensors which will cause a wrong loss calculation, since the output of your model has the shape <code>[batch_size, 1]</code> while the target has the shape <code>[batch_size]</code>.<br>
Since you are using a custom loss function defined as:</p>
<pre><code class="lang-python">def mean_squared_error(ytrue, ypred):
    return torch.mean(((ytrue - ypred) ** 2))
</code></pre>
<p>you would have to make sure the shape checks are used.<br>
Here is a small example:</p>
<pre><code class="lang-python">ytrue = torch.tensor([1.0, 2.0, 3.0])
print(ytrue.shape)
# torch.Size([3])
ypred = torch.tensor([[0.9], [1.9], [2.9]])
print(ypred.shape)
# torch.Size([3, 1])

tmp = (ytrue - ypred)
# this is not what you would expect as the tensors are broadcasted
print(tmp)
# tensor([[ 0.1000,  1.1000,  2.1000],
#         [-0.9000,  0.1000,  1.1000],
#         [-1.9000, -0.9000,  0.1000]])

tmp = (ytrue.unsqueeze(1) - ypred)
# expected results
print(tmp)
# tensor([[0.1000],
#         [0.1000],
#         [0.1000]])
</code></pre>
<p>I would recommend to switch to the built-in loss functions via <code>nn.MSELoss</code> or its functional API <code>F.mse_loss</code>, which would raise a warning:</p>
<pre><code class="lang-python">UserWarning: Using a target size (torch.Size([32])) that is different to the input size (torch.Size([32, 1])). This will likely lead to incorrect results due to broadcasting. Please ensure they have the same size.
</code></pre>
<p>Once this is fixed, the final training loss in the PyTorch model is <code>~0.0027</code> which might be close to the Keras one:</p>
<pre><code class="lang-python">    for i, (x,y) in enumerate(trainloader):
        optimizer.zero_grad()                          
        ypred = torch_model(x)
        loss = F.mse_loss(ypred, y.unsqueeze(1))
        _ = loss.backward()
        _ = optimizer.step()
</code></pre>