Motivation


Two Deep Learning frameworks gather biggest attention - Tensorflow and Pytorch. A lot of Tensorflow popularity among practitioners is due to Keras, which API as of now has been deeply integrated in TF, in the tensorflow.keras module. Two projects - Keras and tensorflow.keras are separate, with first enabling users to change between its backends and second made solely for Tensorflow.

Pytorch is much younger than TF but it has gained a lot of traction since its release. Some time ago, version 1.0.0 was released, indication a big step in the framework maturity.

One of the most important factors when choosing a framework is training and inference speed. While second is of biggest concern in case of model deployment, for researchers and competitive ML practitioners, training time will be the key factor.

So, here comes the speed comparison between current versions of Pytorch and Tensorflow.

Methodology


Comparison is made based on set of reference implementations of various Convolutional Neural Networks architectures. Models for tf.keras are taken from applications module, whereas Pytorch implementations come from eitner torchvision or pretrained-models repository.

Library versions:

  • Tensorflow: 1.12.0
  • Keras: 2.2.4
  • Pytorch: 1.0.0

  • Batch size: model-dependent
  • Image size: 299x299x3
  • Number of samples: 10000
  • Number of epochs: 5
  • Number of runs: 3

Training is performed on a single GTX1080. Model definitions were prepared in a way to be as much similar between frameworks as possible. Dataset was artifically generated - 10k examples in form of a np.array. Final dataset size was thus: [10000, 299, 299, 3]. All models are trained with maximum batch size, that did not cause an error between runs. Batch size was thus equal to minimum batch between two frameworks to have the same one for both TF and Pytorch.

Models

tf.keras:

  • Inception V3
  • ResNet50
  • DenseNet121
  • DenseNet169
  • InceptionResNet V2

Pytorch:

  • Inception V3
  • ResNet50
  • DenseNet121
  • DenseNet169
  • InceptionResNet V2
  • PNASNet
  • NASNet

NASNet in tf.keras was throwing an error, so was finally omitted.

Results

Mean epoch times:

Resulting times present mean epoch time between all runs and epochs for each model in both frameworks.

tf.keras:

TFEpochTimes

Pytorch:

PytorchEpochTimes

Comparison:

ComparisonEpochTimes

First epoch times:

In Tensorflow, first epoch usually takes longer than following ones. This is especially visible in case of more computationally demanding models. In Pytorch, this difference isn’t so pronounced.

FirstEpochTimes

Epoch time excluding the first one:

To give a point of reference, following graph shows mean epoch time, where first epoch was excluded.

NoFirstEpochTimes

Compilation time:

Another aspect, in which TF is slower is the model compilation time. Due to static nature of TF, graph first must be initialized, which results in additional waiting time before the model starts training.

CompilationEpochTimes

Conclusion


Pytorch consistently outperforms tf.keras in terms of model traning speed, when popular CNN architectures are considered. What is more, there is no additional significant increase in training time during the first epoch. Model compilation time is also very quick, which is another argument pro Pytorch. Still, it should be kept in mind that even though tf.keras module contains optimizations prepared especially for Tensorflow, it is still a wrapper over a lower-level DL framework. On the other hand, models benchmarked in Pytorch suffer from no additional overhead.

Code is available on github.

In the future, I hope to benchmark more Pytorch models, such as SE-ResNet, SE-ResNeXt, DPNs & possibly more.