@kpedro88
why don’t you drop the model into lightning and try that (i added the few lines to make that work)?
lightning handles a lot of these performance optimizations such as what you ran into with syncing.
I converted it and ran some quick profiling (all of this took about 10 mins to do end to end).
On CPU was able to get it about 2.3x faster… try on the GPU?

  
      

      GitHub
  

  
    

PyTorchLightning/pytorch-lightning 1

  The lightweight PyTorch wrapper for high-performance AI research. Scale your models, not the boilerplate. - PyTorchLightning/pytorch-lightning


  

  
    
    
  

  



Lightning version
import torch
import torch.nn as nn
from torch.utils.data import TensorDataset, DataLoader, random_split
import torch.optim as optim
from tqdm import tqdm
import numpy as np

# import PL
import pytorch_lightning as pl
pl.seed_everything(0)


class Classifier(pl.LightningModule):

    def __init__(self, model, criterion):
        super().__init__()
        self.model = model
        self.criterion = criterion
 
    def training_step(self, data, batch_idx):
        # train loop
        event_data, target_data = data
        output = self.model(event_data)
        batch_loss = self.criterion(output, target_data)
        self.log('train_loss', batch_loss)
        return batch_loss
    
    def validation_step(self, data, batch_idx):
        # val loop
        event_data, target_data = data
        output = self.model(event_data)
        batch_loss = self.criterion(output, target_data)
        self.log('val_loss', batch_loss)

    def configure_optimizers(self):
        optimizer = optim.Adam(self.model.parameters(), lr=0.001)
        return optimizer

# -------------------
# generate random data
# -------------------
N_train_class0 = N_train_class1 = 1_250_000
event_dim = 8

rng_seed = 0
rng = np.random.Generator(np.random.PCG64(rng_seed))

class0_events = rng.uniform(100,500,size=(N_train_class0, event_dim))
class0_ytarget = np.zeros(shape=(N_train_class0, 1))

class1_events = rng.uniform(100,500,size=(N_train_class1, event_dim))
class1_ytarget = np.zeros(shape=(N_train_class1, 1))

permutation = rng.permutation(N_train_class0 + N_train_class1)
dataset = TensorDataset(
    torch.Tensor(np.concatenate([class0_events, class1_events])[permutation]),
    torch.Tensor(np.concatenate([class0_ytarget, class1_ytarget])[permutation])
)

# prepare for training
pct_valid = 0.2
num_valid = int(len(dataset)*pct_valid)
split_train, split_valid = random_split(dataset, [len(dataset)-num_valid, num_valid])
batch_size = 5000
loader_train = DataLoader(split_train, batch_size=batch_size, shuffle=True)
loader_valid = DataLoader(split_valid, batch_size=batch_size, shuffle=True)


# --------------
# model
# --------------
def init_weights(m):
    if type(m) == nn.Linear:
        torch.nn.init.xavier_uniform(m.weight)
        m.bias.data.fill_(0.0)

hidden_node_counts = [128, 64, 64, 64, 32]
layers = [
    nn.Linear(in_features = event_dim, out_features = hidden_node_counts[0]),
    nn.ReLU()
]
for counter in range(len(hidden_node_counts)-1):
    layers.extend([
        nn.Linear(in_features = hidden_node_counts[counter], out_features = hidden_node_counts[counter+1]),
        nn.ReLU()
    ])
layers.extend([
    nn.Linear(in_features = hidden_node_counts[-1], out_features = 1),
    nn.Sigmoid()
])
model = nn.Sequential(*layers)

print(model)
total_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
print("Trainable params: {}".format(total_params))


criterion = nn.BCELoss()

# -------------
# TRAIN with Lightning
# -------------
classifier = Classifier(model, criterion)
trainer = pl.Trainer(max_epochs=10, gpus=1)
trainer.fit(classifier, loader_train, loader_valid)

I ran it on CPUs (but you can run of gpus by setting gpus=1.
I think your bottleneck also comes from the dataloader… lightning gives you useful warnings when it detects issues.
/Users/williamfalcon/miniconda3/envs/flash/lib/python3.7/site-packages/pytorch_lightning/utilities/distributed.py:49: UserWarning: Your val_dataloader has `shuffle=True`, it is best practice to turn this off for validation and test dataloaders.
  warnings.warn(*args, **kwargs)
/Users/williamfalcon/miniconda3/envs/flash/lib/python3.7/site-packages/pytorch_lightning/utilities/distributed.py:49: UserWarning: The dataloader, val dataloader 0, does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` (try 8 which is the number of cpus on this machine) in the `DataLoader` init to improve performance.
  warnings.warn(*args, **kwargs)
/Users/williamfalcon/miniconda3/envs/flash/lib/python3.7/site-packages/pytorch_lightning/utilities/distributed.py:49: UserWarning: The dataloader, train dataloader, does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` (try 8 which is the number of cpus on this machine) in the `DataLoader` init to improve performance.
  warnings.warn(*args, **kwargs)
Epoch 1:  61%|███████████████████████████████████████████████████▋                                 | 304/500 [00:23<00:15, 12.85it/s, loss=8.94e-12, v_num=4


Profiling
You can also run with profiling to see where the speed issues are coming from.
(profiling might slow things down though…)
Trainer(profiler='pytorch')

---------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  
                                   Name    Self CPU %      Self CPU   CPU total %     CPU total  CPU time avg    # of Calls  
---------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  
                          AddmmBackward        30.23%     763.912ms        30.23%     763.912ms       1.900ms           402  
                               aten::mm        29.64%     749.083ms        29.64%     749.083ms       1.016ms           737  
                            aten::addmm        16.95%     428.369ms        16.95%     428.369ms       1.066ms           402  
                             aten::relu         4.23%     106.824ms         4.23%     106.824ms     318.878us           335  
                        aten::threshold         4.02%     101.632ms         4.02%     101.632ms     303.378us           335  
                          ReluBackward0         3.97%     100.330ms         3.97%     100.330ms     299.493us           335  
               aten::threshold_backward         3.89%      98.345ms         3.89%      98.345ms     293.566us           335  
                              aten::sum         1.61%      40.591ms         1.61%      40.591ms      86.548us           469  
                            aten::copy_         1.47%      37.197ms         1.47%      37.197ms      46.265us           804  
                                aten::t         0.78%      19.632ms         0.78%      19.632ms      10.104us          1943  
             aten::binary_cross_entropy         0.36%       9.013ms         0.36%       9.013ms     134.518us            67  
        torch::autograd::AccumulateGrad         0.32%       8.046ms         0.32%       8.046ms      10.008us           804  
                        aten::transpose         0.30%       7.688ms         0.30%       7.688ms       3.957us          1943  
                            aten::empty         0.26%       6.545ms         0.26%       6.545ms       2.505us          2613  
                          aten::sigmoid         0.22%       5.542ms         0.22%       5.542ms      41.361us           134  
                             aten::add_         0.20%       5.042ms         0.20%       5.042ms       6.367us           792  
                       aten::as_strided         0.17%       4.377ms         0.17%       4.377ms       1.519us          2881  
                              TBackward         0.14%       3.533ms         0.14%       3.533ms       8.788us           402  
                             aten::view         0.14%       3.463ms         0.14%       3.463ms       8.613us           402  
                             aten::mean         0.11%       2.833ms         0.11%       2.833ms      42.280us            67  
---------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  
Self CPU time total: 2.527s

Or try any of the other profilers… here we use the “simple” profiler
Trainer(profiler='simple')

Profiler Report

Action                                  |  Mean duration (s)    |Num calls              |  Total time (s)       |  Percentage %         |
------------------------------------------------------------------------------------------------------------------------------------
Total                                   |  -                    |_                      |  37.739               |  100 %                |
------------------------------------------------------------------------------------------------------------------------------------
run_training_epoch                      |  37.493               |1                      |  37.493               |  99.348               |
get_train_batch                         |  0.053623             |400                    |  21.449               |  56.836               |
run_training_batch                      |  0.024157             |400                    |  9.6627               |  25.604               |
optimizer_step_and_closure_0            |  0.023673             |400                    |  9.4692               |  25.091               |
training_step_and_backward              |  0.022778             |400                    |  9.1112               |  24.143               |
model_backward                          |  0.013783             |400                    |  5.5132               |  14.609               |
model_forward                           |  0.0088057            |400                    |  3.5223               |  9.3333               |
training_step                           |  0.0085555            |400                    |  3.4222               |  9.068                |
evaluation_step_and_end                 |  0.0089042            |102                    |  0.90823              |  2.4066               |
validation_step                         |  0.0087426            |102                    |  0.89174              |  2.3629               |
on_train_batch_end                      |  0.00036939           |400                    |  0.14776              |  0.39152              |
cache_result                            |  1.859e-05            |2323                   |  0.043184             |  0.11443              |
on_validation_batch_end                 |  0.0002122            |102                    |  0.021644             |  0.057353             |
on_batch_start                          |  3.5443e-05           |400                    |  0.014177             |  0.037566             |
on_after_backward                       |  1.7219e-05           |400                    |  0.0068876            |  0.018251             |
on_batch_end                            |  1.3641e-05           |400                    |  0.0054563            |  0.014458             |
on_before_zero_grad                     |  1.2449e-05           |400                    |  0.0049796            |  0.013195             |
on_train_batch_start                    |  9.6932e-06           |400                    |  0.0038773            |  0.010274             |
training_step_end                       |  8.5934e-06           |400                    |  0.0034374            |  0.0091083            |
on_validation_end                       |  0.0016228            |2                      |  0.0032456            |  0.0086001            |
on_validation_batch_start               |  2.1877e-05           |102                    |  0.0022315            |  0.005913             |
validation_step_end                     |  8.1364e-06           |102                    |  0.00082992           |  0.0021991            |
on_train_end                            |  0.00024225           |1                      |  0.00024225           |  0.00064191           |
on_epoch_start                          |  0.00021004           |1                      |  0.00021004           |  0.00055657           |
on_validation_start                     |  7.725e-05            |2                      |  0.0001545            |  0.00040939           |
on_train_start                          |  0.00014138           |1                      |  0.00014138           |  0.00037461           |
on_validation_epoch_end                 |  1.4126e-05           |2                      |  2.8251e-05           |  7.4859e-05           |
on_epoch_end                            |  1.7333e-05           |1                      |  1.7333e-05           |  4.5929e-05           |
on_validation_epoch_start               |  7.854e-06            |2                      |  1.5708e-05           |  4.1623e-05           |
on_before_accelerator_backend_setup     |  1.2458e-05           |1                      |  1.2458e-05           |  3.3011e-05           |
on_train_epoch_end                      |  9.25e-06             |1                      |  9.25e-06             |  2.4511e-05           |
on_fit_start                            |  8.709e-06            |1                      |  8.709e-06            |  2.3077e-05           |
on_train_epoch_start                    |  8e-06                |1                      |  8e-06                |  2.1198e-05           |


logging
And finally, if you want more speed ups, turn off the logger (lightning automatically generates tensorboard files).
Trainer(logger=False)


Bonus
Notice that Lightning is much simpler (but very different than keras).
If you want more of the “Keras” feel in pytorch, use our other library Flash (built on top of lightning)
Just add these last lines to your code
import Flash

# task
classifier = flash.Task(model, loss_fn=criterion, optimizer=optimizer)

# train
flash.Trainer(max_epochs=10, gpus=1).fit(classifier, loader_train, loader_valid)


Example (with num_workers)
notice on the profiler that you spent a non trivial amount of time on data loading and lightning
gave you a warning about num_workers?
Dataloader(num_workers=8)

Now if I run and profile again… we eliminate the data loading issue and you are left with the model optimization speed ups (which lightning handles most of the bottlenecks there)
Profiler Report

Action                                  |  Mean duration (s)    |Num calls              |  Total time (s)       |  Percentage %         |
------------------------------------------------------------------------------------------------------------------------------------
Total                                   |  -                    |_                      |  16.307               |  100 %                |
------------------------------------------------------------------------------------------------------------------------------------
run_training_epoch                      |  15.476               |1                      |  15.476               |  94.909               |
run_training_batch                      |  0.029962             |400                    |  11.985               |  73.496               |
optimizer_step_and_closure_0            |  0.029476             |400                    |  11.791               |  72.305               |
training_step_and_backward              |  0.028529             |400                    |  11.412               |  69.983               |
model_backward                          |  0.015577             |400                    |  6.2306               |  38.209               |
model_forward                           |  0.012756             |400                    |  5.1025               |  31.291               |
training_step                           |  0.012475             |400                    |  4.9902               |  30.602               |
evaluation_step_and_end                 |  0.017031             |102                    |  1.7372               |  10.653               |
validation_step                         |  0.016785             |102                    |  1.712                |  10.499               |
get_train_batch                         |  0.0024146            |400                    |  0.96584              |  5.923                |
on_train_batch_end                      |  0.00035222           |400                    |  0.14089              |  0.86398              |
cache_result                            |  2.2822e-05           |2323                   |  0.053016             |  0.32512              |
on_validation_batch_end                 |  0.00017669           |102                    |  0.018023             |  0.11052              |
on_batch_start                          |  2.3768e-05           |400                    |  0.0095072            |  0.058303             |
on_after_backward                       |  1.6999e-05           |400                    |  0.0067996            |  0.041698             |
on_batch_end                            |  1.5546e-05           |400                    |  0.0062184            |  0.038134             |
on_before_zero_grad                     |  1.4941e-05           |400                    |  0.0059764            |  0.03665              |
on_validation_end                       |  0.0021403            |2                      |  0.0042806            |  0.026251             |
training_step_end                       |  1.0234e-05           |400                    |  0.0040936            |  0.025104             |
on_train_batch_start                    |  9.4705e-06           |400                    |  0.0037882            |  0.023231             |
on_validation_batch_start               |  1.853e-05            |102                    |  0.00189              |  0.011591             |
validation_step_end                     |  1.3889e-05           |102                    |  0.0014167            |  0.0086879            |
on_validation_start                     |  0.00023427           |2                      |  0.00046854           |  0.0028733            |
on_validation_epoch_end                 |  0.00011319           |2                      |  0.00022638           |  0.0013882            |
on_epoch_start                          |  0.00022617           |1                      |  0.00022617           |  0.001387             |
on_train_end                            |  0.00022525           |1                      |  0.00022525           |  0.0013813            |
on_train_start                          |  0.000178             |1                      |  0.000178             |  0.0010916            |
on_epoch_end                            |  1.8708e-05           |1                      |  1.8708e-05           |  0.00011473           |
on_validation_epoch_start               |  9.2705e-06           |2                      |  1.8541e-05           |  0.0001137            |
on_train_epoch_start                    |  1.0667e-05           |1                      |  1.0667e-05           |  6.5415e-05           |
on_train_epoch_end                      |  9.209e-06            |1                      |  9.209e-06            |  5.6474e-05           |
on_fit_start                            |  8.209e-06            |1                      |  8.209e-06            |  5.0341e-05           |
on_before_accelerator_backend_setup     |  6.417e-06            |1                      |  6.417e-06            |  3.9352e-05           |


summary
so, in summary… it just looks like keras does more optimizations under the hood to help the user avoid mistakes. This is also something that lightning does.

Bonus 2
Also notice lightning gives this warning…
/Users/williamfalcon/miniconda3/envs/flash/lib/python3.7/site-packages/pytorch_lightning/utilities/distributed.py:49: UserWarning: Your val_dataloader has `shuffle=True`, it is best practice to turn this off for validation and test dataloaders.

it caught a mistake you made… don’t shuffle your validation data (it’s bad practice)