Add: Enter training time for BF16 using AVX512

Update IntelPyTorch training notebook to include environment variable settings for AVX512 and BF16 optimizations

Signed-off-by: Francisco R Castro G <[email protected]>
Signed-off-by: jafraustro <[email protected]>

Author: jafraustro

AI-and-Analytics/Features-and-Functionality/IntelPyTorch_TrainingOptimizations_AMX_BF16/IntelPyTorch_TrainingOptimizations_AMX_BF16.ipynb

@@ -254,29 +254,20 @@
    "outputs": [],
    "source": [
     "print(\"Training model with FP32\")\n",
+    "os.environ[\"ONEDNN_MAX_CPU_ISA\"] = \"AVX512_CORE_AMX\"\n",
     "fp32_training_time = trainModel(train_loader, modelName=\"fp32\", dataType=\"fp32\")"
    ]
   },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "75aafe25-4f7d-42ad-92ed-3438bd78c00b",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "print(\"Training model with BF16 with AVX512\")\n",
-    "!python pytorch_training_avx512_bf16.py"
-   ]
-  },
   {
    "cell_type": "code",
    "execution_count": null,
    "id": "a9bd6dec",
    "metadata": {},
    "outputs": [],
    "source": [
-    "# Record the training time for BF16 using AVX512\n",
-    "bf16_avx512_training_time = None   #TODO: enter in training time"
+    "print(\"Training model with BF16 with AVX512\")\n",
+    "os.environ[\"ONEDNN_MAX_CPU_ISA\"] = \"AVX512_CORE_BF16\"\n",
+    "bf16_avx512_training_time = trainModel(train_loader, modelName=\"bf16_noAmx\", dataType=\"bf16\")"
    ]
   },
   {
@@ -287,6 +278,7 @@
    "outputs": [],
    "source": [
     "print(\"Training model with BF16 with Intel® AMX\")\n",
+    "os.environ[\"ONEDNN_MAX_CPU_ISA\"] = \"AVX512_CORE_AMX\"\n",
     "bf16_amx_training_time = trainModel(train_loader, modelName=\"bf16_withAmx\", dataType=\"bf16\")"
    ]
   },
@@ -391,9 +383,9 @@
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "Python 3.8.10 64-bit (microsoft store)",
+   "display_name": "pytorch",
    "language": "python",
-   "name": "python3"
+   "name": "pytorch"
   },
   "language_info": {
    "codemirror_mode": {
@@ -405,7 +397,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.8.10"
+   "version": "3.11.0"
   },
   "vscode": {
    "interpreter": {