Performance of Pytorch vs pure Torch

[adrianomartinelli]

Hi! I very much like the lightning library and recently started to experiment with it. However, I observed that the lightning training takes significantly longer than pure torch (I also saw that you just opened an issue (#12398 @carmocca) related to performance yourself). I performed a little benchmark and wondered if these results are expected and if performance differences are less on gpu.

Benchmark
I tested two model sizes (small model with a few thousand parameters and a large one with millions) and compared the duration of training the models for 10 or 250 epochs on a large data set (1e6 observations) and a small data set (1e3 observations) with a batch size of 256. I tried to disable all lightning functionalities that may slowdown training that I am aware of (like checkpoints or progress bars). The machine was only used for the benchmark (see machine_info.txt for details).

Observations:

• Lighting training takes ~60% longer for the small model, seems not to be influenced by number of epochs or size of data set
• Lightning training takes ~25% longer for the larger model, seems not to be influenced by number of epochs or size of data set

Thanks!
Adriano

Supporting Material
benchmark.txt: output of running benchmark.py.txt
machine_info.txt: information about the machine used to run the benchmark

# %%
import torch
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader

import pytorch_lightning as pl
from pytorch_lightning.utilities.types import STEP_OUTPUT
from typing import Optional

import sys
import time
from pathlib import Path

p = Path('~/benchmark.out').expanduser()

# %%
mode = sys.argv[1]
data_size = int(sys.argv[2])
model_size = sys.argv[3]
epochs = int(sys.argv[4])
batch_size = int(sys.argv[5])
N = int(sys.argv[6])
print(f'{data_size}\t{mode}\t{model_size}\t{epochs}\t{batch_size}\t{N}')
# %% data
X = torch.randn(data_size, 50)


class DataSet(Dataset):
    def __init__(self, X):
        super().__init__()
        self.data = X

    def __getitem__(self, item):
        return self.data[item], self.data[item]

    def __len__(self):
        return len(self.data)


# %% models

def small_model(d1):
    class Model(nn.Module):

        def __init__(self):
            super().__init__()
            self.model = nn.Sequential(nn.Linear(d1, 10), nn.ReLU(),
                                       nn.Linear(10, 2), nn.ReLU(),
                                       nn.Linear(2, 10), nn.ReLU(),
                                       nn.Linear(10, d1), nn.Sigmoid())

        def forward(self, x):
            return self.model(x)

    return Model()


def large_model(d1):
    class Model(nn.Module):

        def __init__(self):
            super().__init__()
            self.model = nn.Sequential(nn.Linear(d1, 50), nn.ReLU(),
                                       nn.Linear(50, 50), nn.ReLU(),
                                       nn.Linear(50, 50), nn.ReLU(),
                                       nn.Linear(50, 50), nn.ReLU(),
                                       nn.Linear(50, 2), nn.ReLU(),
                                       nn.Linear(2, 50), nn.ReLU(),
                                       nn.Linear(50, 50), nn.ReLU(),
                                       nn.Linear(50, 50), nn.ReLU(),
                                       nn.Linear(50, 50), nn.ReLU(),
                                       nn.Linear(50, d1), nn.Sigmoid())

        def forward(self, x):
            return self.model(x)

    return Model()


def lightning_module(model):
    class Model(pl.LightningModule):
        def __init__(self):
            super().__init__()
            self.loss = nn.MSELoss()
            self.doLog = False
            self.learning_rate = 1e-3
            self.model = model

        def forward(self, x):
            return self.model(x)

        def training_step(self, batch, batch_idx) -> STEP_OUTPUT:
            x, y = batch
            yhat = self.model(x)
            loss = self.loss(yhat, y)
            if self.doLog:
                self.log('test_loss', loss.detach())
                self.log_metrics('train', y, yhat)
            return loss

        def validation_step(self, batch, batch_idx) -> Optional[STEP_OUTPUT]:
            x, y = batch
            yhat = self.model(x)
            loss = self.loss(yhat, y)
            if self.doLog:
                self.log('test_loss', loss.detach())
                self.log_metrics('validate', y, yhat)
            return loss

        def configure_optimizers(self):
            # https://pytorch-lightning.readthedocs.io/en/latest/api/pytorch_lightning.core.lightning.html#pytorch_lightning.core.lightning.LightningModule.configure_optimizers
            optimizer = torch.optim.Adam(self.parameters(), lr=self.learning_rate)
            return optimizer

        def log_metrics(self, step, y, yhat):
            # https://pytorch-lightning.readthedocs.io/en/stable/extensions/logging.html

            for metric in self.metric_attrs:
                m = getattr(self, metric)
                m(yhat, y)
                self.log(f'{step}_{metric}', m)

    return Model()


# %%

def fit_torch(model, train_dataloader, epochs=10):
    def train_loop(dataloader, model, loss_fn, optimizer):
        for batch, (X, y) in enumerate(dataloader):
            # Compute prediction and loss
            pred = model(X)
            loss = loss_fn(pred, y)

            # Backpropagation
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

    optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)

    for t in range(epochs):
        train_loop(train_dataloader, model, nn.MSELoss(), optimizer)


def fit_lightning(model, train_dataloader, epochs=10):
    trainer = pl.Trainer(max_epochs=epochs, logger=False, enable_checkpointing=False, enable_progress_bar=False,
                         enable_model_summary=False)
    trainer.fit(model, train_dataloaders=train_dataloader)


# %%

ds = DataSet(X)

d1 = X.shape[-1]
model = small_model(d1) if model_size == 'small' else large_model(d1)
model = lightning_module(model) if mode == 'light' else model

res = []
for i in range(N):
    fit = fit_torch if mode == 'torch' else fit_lightning
    train_dataloader = DataLoader(ds, batch_size=batch_size, num_workers=0)

    start = time.time()
    fit(model, train_dataloader, epochs=epochs)
    res.append((time.time() - start))

res = torch.tensor(res).float()
# %%

with open(p, 'a') as f:
    s = f'{data_size}\t{mode}\t{model_size}\t{epochs}\t{batch_size}\t{N}\t{res.mean():.2f}\t{res.max():.2f}\t{res.min():.2f}\n'
    f.write(s)

[carmocca]

Can you also share the terminal calls for each run?

[adrianomartinelli]

I executed this bash script from the terminal (changed the order of the calls in the meantime).

run_benchmark.sh.txt

#!/bin/bash

#conda activate starProtocols
#
python benchmark.py 'torch' 10000 small 10 256 10
python benchmark.py 'torch' 10000 large 10 256 10
python benchmark.py 'light' 10000 small 10 256 10
python benchmark.py 'light' 10000 large 10 256 10
#
python benchmark.py 'torch' 10000 small 250 256 10
python benchmark.py 'torch' 10000 large 250 256 10
python benchmark.py 'light' 10000 small 250 256 10
python benchmark.py 'light' 10000 large 250 256 10
#
python benchmark.py 'torch' 1000000 small 10 256 10
python benchmark.py 'torch' 1000000 large 10 256 10
python benchmark.py 'light' 1000000 small 10 256 10
python benchmark.py 'light' 1000000 large 10 256 10

[rohitgr7]

I tried running your benchmark script and noticed that the experiments run only for a few seconds per iteration (<5 seconds). IMO, this might not be an ideal case to compare scripts where the difference will come out to be (0.5-1 sec) but the % might look huge. To make it comparable I used 100 layers in your large model and found out that the average time per iteration in both cases was similar (~23.5 sec)