Keras vs Pytorch_lightning execution time

[sushmit86]

I am trying to convert a Keras code to Pytorch_lightning, I see a significant difference between the two. Keras is running much faster than PL. Here is my code below. Here is also a Colab link for the same Colab link. Please note I have tested both with CPU and GPU. In both cases Keras is running much faster.

## Import Libraries
import pandas as pd
from keras.datasets import boston_housing
import numpy as np
import torch
import pytorch_lightning as pl
from torch.utils.data import DataLoader,TensorDataset
from torch.nn import functional as F
from pytorch_lightning.loggers import CSVLogger
import torchmetrics
from torch import nn
from keras import models
from keras import layers
## Keras code
k = 4
num_val_samples = len(train_data) // k
num_epochs = 100
all_scores = []
def build_model():
  model = models.Sequential()
  model.add(layers.Dense(64, activation='relu',
  input_shape=(train_data.shape[1],)))
  model.add(layers.Dense(64, activation='relu'))
  model.add(layers.Dense(1))
  model.compile(optimizer='rmsprop', loss='mse', metrics=['mae'])
  return model
for i in range(k):
  print('processing fold #', i)
  val_data = train_data[i * num_val_samples: (i + 1) * num_val_samples]
  val_targets = train_targets[i * num_val_samples: (i + 1) * num_val_samples]
  partial_train_data = np.concatenate(
  [train_data[:i * num_val_samples],
  train_data[(i + 1) * num_val_samples:]],
  axis=0)
  partial_train_targets = np.concatenate(
  [train_targets[:i * num_val_samples],
  train_targets[(i + 1) * num_val_samples:]],
  axis=0)
  model = build_model()
  model.fit(partial_train_data, partial_train_targets,
  epochs=num_epochs, batch_size=1, verbose=0)
  val_mse, val_mae = model.evaluate(val_data, val_targets, verbose=0)
all_scores.append(val_mae)
##Pytorch_lightning code
class boston_reg(pl.LightningModule):
  def __init__(self,net):
    super().__init__()
    self.net = net
    self.mae = torchmetrics.MeanAbsoluteError()
  def forward(self,x):
    x = self.net(x)
    return x
  def loss_fn(self,y_hat,y):
    loss = nn.MSELoss()
    return loss(y_hat,y)
  def training_step(self,train_batch,batch_idx):
    x,y = train_batch
    y = y.view(-1,1)
    y_hat = self.forward(x)
    loss = self.loss_fn(y_hat,y)
    self.log('train_loss',loss, on_step = False, on_epoch = True)
    mae = self.mae(y_hat,y)
    self.log('train_mae',mae, on_step = False, on_epoch = True)
    return loss
  def test_step(self,test_batch,batch_idx):
    x,y = test_batch
    y = y.view(-1,1)
    y_hat = self(x)
    loss = self.loss_fn(y_hat,y)
    self.log('test_loss',loss)
    mae = self.mae(y_hat,y)
    self.log('test_mae',mae)
#   def validation_step(self,train_batch,batch_idx):
#     x,y = train_batch
#     y_hat = self(x)
#     loss = self.loss_fn(y_hat,y)
#     self.log('val_loss',loss, on_step = False, on_epch = True)
    
  def configure_optimizers(self):
    optimizer = torch.optim.RMSprop(self.parameters(), lr=0.001,alpha=0.9)
    #optimizer = torch.optim.SGD(self.parameters(), lr=0.001)
    return optimizer

class Light_DataModule_boston(pl.LightningDataModule):
  def __init__(self,x_train,y_train,x_test,y_test):
    '''
    x_train: Torch tensors
    y_train: Torch tensors
    
    '''
    super().__init__()
    self.x_val = x_test
    self.partial_x_train = x_train
    self.y_val = y_test
    self.partial_y_train = y_train
    self.x_test = x_test
    self.y_test = y_test

  def train_dataloader(self):
    return DataLoader(TensorDataset(self.partial_x_train,self.partial_y_train), batch_size=1,shuffle=True,num_workers=1)
#   def val_dataloader(self):
#     return DataLoader(TensorDataset(self.x_val,self.y_val), batch_size=512,shuffle=False,num_workers=1)
  def test_dataloader(self):
    return DataLoader(TensorDataset(self.x_test,self.y_test), batch_size=1,shuffle=False,num_workers=1)
k = 4
num_val_samples = len(train_data) // k
num_epochs = 100
all_scores = []
torch_x_train = torch.Tensor(train_data)
torch_y_train = torch.Tensor(train_targets)
len(torch_x_train)
all_scores = []
num_epochs = 100
# net = nn.Sequential(nn.Linear(torch_x_train.shape[1], 64),nn.ReLU(),nn.Linear(64,64),nn.ReLU(),nn.Linear(64,1))
# logger = CSVLogger("logs", name="boston_reg",version = "version_0")
# model  = boston_reg(net)
for i in range(k):
  print('processing fold #', i)
  net = nn.Sequential(nn.Linear(torch_x_train.shape[1], 64),nn.ReLU(),nn.Linear(64,64),nn.ReLU(),nn.Linear(64,1))
  logger = CSVLogger("logs", name="boston_reg",version = "version_0")
  model  = boston_reg(net)
  val_data = torch_x_train[i * num_val_samples: (i + 1) * num_val_samples]
  val_targets = torch_y_train[i * num_val_samples: (i + 1) * num_val_samples]
  partial_train_data = torch.cat([torch_x_train[:i * num_val_samples],torch_x_train[(i + 1) * num_val_samples:]],axis=0)
  partial_train_targets = torch.cat([torch_y_train[:i * num_val_samples],torch_y_train[(i + 1) * num_val_samples:]],axis=0)
  data_module = Light_DataModule_boston(partial_train_data,partial_train_targets,val_data,val_targets)
  trainer = pl.Trainer(max_epochs=100,gpus=0,progress_bar_refresh_rate = 0,logger = logger)
  trainer.fit(model, data_module)
  test_results = trainer.test(model,verbose = True)
  all_scores.append(test_results[0]['test_mae'])
print(all_scores)

[tchaton]

Dear @sushmit86,

Have you tried for a more significant model size ?

Side note: You could use Trainer(profiler="simple") to find out where the time is being spent on the Lightning side.

Best,
T.C

[sushmit86]

Thanks @tchaton for this here is the profiler report. Any ideas

FIT Profiler Report

Action                             	|  Mean duration (s)	|Num calls      	|  Total time (s) 	|  Percentage %   	|
--------------------------------------------------------------------------------------------------------------------------------------
Total                              	|  -              	|_              	|  161.62         	|  100 %          	|
--------------------------------------------------------------------------------------------------------------------------------------
run_training_epoch                 	|  1.6159         	|100            	|  161.59         	|  99.98          	|
run_training_batch                 	|  0.003105       	|30300          	|  94.082         	|  58.212         	|
optimizer_step_and_closure_0       	|  0.0025525      	|30300          	|  77.342         	|  47.854         	|
training_step_and_backward         	|  0.0019799      	|30300          	|  59.991         	|  37.119         	|
get_train_batch                    	|  0.001624       	|30300          	|  49.208         	|  30.447         	|
model_forward                      	|  0.0013005      	|30300          	|  39.407         	|  24.382         	|
training_step                      	|  0.0011345      	|30300          	|  34.375         	|  21.269         	|
backward                           	|  0.00042349     	|30300          	|  12.832         	|  7.9394         	|
training_batch_to_device           	|  0.00017718     	|30300          	|  5.3686         	|  3.3218         	|
on_train_batch_end                 	|  5.0194e-05     	|30300          	|  1.5209         	|  0.94103        	|
on_batch_start                     	|  3.2736e-05     	|30300          	|  0.9919         	|  0.61372        	|
on_train_batch_start               	|  2.7406e-05     	|30300          	|  0.8304         	|  0.5138         	|
on_before_optimizer_step           	|  2.6072e-05     	|30300          	|  0.78998        	|  0.48879        	|
on_after_backward                  	|  2.367e-05      	|30300          	|  0.7172         	|  0.44376        	|
on_batch_end                       	|  2.3662e-05     	|30300          	|  0.71695        	|  0.4436         	|
on_before_zero_grad                	|  2.318e-05      	|30300          	|  0.70235        	|  0.43457        	|
on_before_backward                 	|  2.0433e-05     	|30300          	|  0.61911        	|  0.38307        	|
training_step_end                  	|  2.01e-05       	|30300          	|  0.60902        	|  0.37682        	|
on_train_epoch_end                 	|  0.0025426      	|100            	|  0.25426        	|  0.15732        	|
on_epoch_start                     	|  2.7504e-05     	|100            	|  0.0027504      	|  0.0017018      	|
on_epoch_end                       	|  2.5981e-05     	|100            	|  0.0025981      	|  0.0016075      	|
on_train_epoch_start               	|  1.9341e-05     	|100            	|  0.0019341      	|  0.0011967      	|
on_fit_start                       	|  5.2976e-05     	|1              	|  5.2976e-05     	|  3.2778e-05     	|
on_train_start                     	|  5.0506e-05     	|1              	|  5.0506e-05     	|  3.125e-05      	|
on_train_end                       	|  4.703e-05      	|1              	|  4.703e-05      	|  2.9099e-05     	|
on_train_dataloader                	|  2.8969e-05     	|1              	|  2.8969e-05     	|  1.7924e-05     	|
on_before_accelerator_backend_setup	|  1.5599e-05     	|1              	|  1.5599e-05     	|  9.6517e-06  

[tchaton]

You could try to increase the number of workers. It seems the loading of data takes 30 % of the training which is quite high or the batch_size which is at 1.

[sushmit86]

I have already tried with a different sets of num of works . The batch size is the same in both cases.