JAX Tutorials: Small changes in speed comparison

xgxg1314 · Jun 8, 2022 · 897b6af · 897b6af
1 parent a5a7f0d
commit 897b6af
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 9 deletions.
diff --git a/docs/tutorial_notebooks/JAX/tutorial5/Inception_ResNet_DenseNet.ipynb b/docs/tutorial_notebooks/JAX/tutorial5/Inception_ResNet_DenseNet.ipynb
@@ -36,7 +36,7 @@
    "source": [
     "---\n",
     "\n",
-    "**Speed comparison**: We note the training times for all models in the PyTorch and the JAX implementation below (PyTorch v1.11, JAX v0.3.13). The models were trained on the same hardware (NVIDIA RTX3090, 24 core CPU) and we slightly adjusted the tutorials to use the exact same training settings (200 epochs, data loading parameters, evaluation schedule, etc.). Overall, the JAX implementation required 60-70% less training time!\n",
+    "**Speed comparison**: We note the training times for all models in the PyTorch and the JAX implementation below (PyTorch v1.11, JAX v0.3.13). The models were trained on the same hardware (NVIDIA RTX3090, 24 core CPU) and we slightly adjusted the tutorials to use the exact same training settings (200 epochs, data loading parameters, evaluation schedule, etc.). Overall, the JAX implementation is about *2.5-3.4x faster* than PyTorch!\n",
     "    \n",
     "| Models                |   PyTorch   |     JAX     |\n",
     "|-----------------------|:-----------:|:-----------:|\n",
@@ -2030,7 +2030,7 @@
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "Python 3",
+   "display_name": "Python 3 (ipykernel)",
    "language": "python",
    "name": "python3"
   },
@@ -2044,7 +2044,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.7.9"
+   "version": "3.8.2"
   }
  },
  "nbformat": 4,

diff --git a/docs/tutorial_notebooks/JAX/tutorial6/Transformers_and_MHAttention.ipynb b/docs/tutorial_notebooks/JAX/tutorial6/Transformers_and_MHAttention.ipynb
@@ -35,7 +35,7 @@
    "source": [
     "---\n",
     "\n",
-    "**Speed comparison**: We note the training times for all models in the PyTorch and the JAX implementation below (PyTorch v1.11, JAX v0.3.13). The models were trained on the same hardware (NVIDIA RTX3090, 24 core CPU) and we slightly adjusted the tutorials to use the exact same training settings (same data loading parameters, evaluation schedule, etc.). Overall, the JAX implementation required ~75% less training time!\n",
+    "**Speed comparison**: We note the training times for all models in the PyTorch and the JAX implementation below (PyTorch v1.11, JAX v0.3.13). The models were trained on the same hardware (NVIDIA RTX3090, 24 core CPU) and we slightly adjusted the tutorials to use the exact same training settings (same data loading parameters, evaluation schedule, etc.). Overall, the JAX implementation is almost *4x faster* than PyTorch!\n",
     "    \n",
     "| Models            |   PyTorch   |     JAX    |\n",
     "|-------------------|:-----------:|:----------:|\n",
@@ -20343,7 +20343,7 @@
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "Python 3",
+   "display_name": "Python 3 (ipykernel)",
    "language": "python",
    "name": "python3"
   },
@@ -20357,7 +20357,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.7.9"
+   "version": "3.8.2"
   }
  },
  "nbformat": 4,

diff --git a/docs/tutorial_notebooks/JAX/tutorial9/AE_CIFAR10.ipynb b/docs/tutorial_notebooks/JAX/tutorial9/AE_CIFAR10.ipynb
@@ -35,7 +35,7 @@
    "source": [
     "---\n",
     "\n",
-    "**Speed comparison**: We note the training times for all models in the PyTorch and the JAX implementation below (PyTorch v1.11, JAX v0.3.13). The models were trained on the same hardware (NVIDIA RTX3090, 24 core CPU) and we slightly adjusted the tutorials to use the exact same training settings (same data loading parameters, evaluation schedule, etc.). Overall, the JAX implementation required ~40% less training time!\n",
+    "**Speed comparison**: We note the training times for all models in the PyTorch and the JAX implementation below (PyTorch v1.11, JAX v0.3.13). The models were trained on the same hardware (NVIDIA RTX3090, 24 core CPU) and we slightly adjusted the tutorials to use the exact same training settings (same data loading parameters, evaluation schedule, etc.). Overall, the JAX implementation is about *1.7x faster* than PyTorch!\n",
     "    \n",
     "| Models           |   PyTorch   |     JAX    |\n",
     "|------------------|:-----------:|:----------:|\n",
@@ -7285,7 +7285,7 @@
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "Python 3",
+   "display_name": "Python 3 (ipykernel)",
    "language": "python",
    "name": "python3"
   },
@@ -7299,7 +7299,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.7.9"
+   "version": "3.8.2"
   }
  },
  "nbformat": 4,