Roll back to use python script.

- BF16 without Intel® AMX example is executed from a python script. The bf16_noAmx_training_time variables is saved in a txt file.

- Use torch.amp.autocast('cpu') instead of torch.cpu.amp.autocast()

Signed-off-by: jafraustro <[email protected]>

Author: jafraustro

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

@@ -118,7 +118,7 @@
     "## Training the Model\n",
     "The function trainModel() will train the Resnet50 model based on the whether Intel® AMX should be enabled, and whether to use FP32 or BF16 data type. The environment variable `ONEDNN_MAX_CPU_ISA` is used to enable or disable Intel® AMX. **Note that this environment variable is only initialized once.** This means to run with Intel® AMX and VNNI, there will need to be separate processes. The best practice is to set this environment variable before running your script. For more information, refer to the [oneDNN documentation on CPU Dispatcher Control](https://www.intel.com/content/www/us/en/develop/documentation/onednn-developer-guide-and-reference/top/performance-profiling-and-inspection/cpu-dispatcher-control.html). \n",
     "\n",
-    "To use BF16 in operations, use the `torch.cpu.amp.autocast()` function to perform forward and backward propagation."
+    "To use BF16 in operations, use the `torch.amp.autocast('cpu')` function to perform forward and backward propagation."
    ]
   },
   {
@@ -128,7 +128,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "os.environ[\"ONEDNN_MAX_CPU_ISA\"] = \"AVX512_CORE_BF16\""
+    "os.environ[\"ONEDNN_MAX_CPU_ISA\"] = \"AVX512_CORE_AMX\""
    ]
   },
   {
@@ -171,7 +171,7 @@
     "    for batch_idx, (data, target) in enumerate(train_loader):\n",
     "        optimizer.zero_grad()\n",
     "        if \"bf16\" == dataType:\n",
-    "            with torch.cpu.amp.autocast():   # Auto Mixed Precision\n",
+    "            with torch.amp.autocast('cpu'):   # Auto Mixed Precision\n",
     "                # Setting memory_format to torch.channels_last could improve performance with 4D input data. This is optional.\n",
     "                data = data.to(memory_format=torch.channels_last)\n",
     "                output = model(data)\n",
@@ -240,8 +240,8 @@
     "## Training with FP32 and BF16, including Intel® AMX\n",
     "Train the Resnet50 model in three different cases:\n",
     "1. FP32 (baseline)  \n",
-    "2. BF16 without Intel® AMX  \n",
-    "3. BF16 with Intel® AMX  \n",
+    "2. BF16 with Intel® AMX\n",
+    "x. BF16 without Intel® AMX\n",
     "\n",
     "The training time is recorded."
    ]
@@ -254,32 +254,33 @@
    "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": "a9bd6dec",
+   "id": "3faaf5de",
    "metadata": {},
    "outputs": [],
    "source": [
-    "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\")"
+    "print(\"Training model with BF16 with Intel® AMX\")\n",
+    "bf16_amx_training_time = trainModel(train_loader, modelName=\"bf16_withAmx\", dataType=\"bf16\")"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "2fdc8a70-509a-4714-8524-084f34e287c3",
+   "id": "a9bd6dec",
    "metadata": {},
    "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\")"
+    "print(\"Training model with BF16 with AVX512\")\n",
+    "!python pytorch_training_avx512_bf16.py\n",
+    "\n",
+    "# Read the variable\n",
+    "with open('bf16_avx512_training_time.txt', 'r') as f:\n",
+    "    bf16_avx512_training_time = float(f.read().strip())"
    ]
   },
   {
@@ -383,9 +384,9 @@
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "pytorch",
+   "display_name": "pytorch_test",
    "language": "python",
-   "name": "pytorch"
+   "name": "pytorch_test"
   },
   "language_info": {
    "codemirror_mode": {
@@ -398,11 +399,6 @@
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
    "version": "3.11.0"
-  },
-  "vscode": {
-   "interpreter": {
-    "hash": "ed6ae0d06e7bec0fef5f1fb38f177ceea45508ce95c68ed2f49461dd6a888a39"
-   }
   }
  },
  "nbformat": 4,

AI-and-Analytics/Features-and-Functionality/IntelPyTorch_TrainingOptimizations_AMX_BF16/pytorch_training_avx512_bf16.py

@@ -56,7 +56,7 @@ def trainModel(train_loader, modelName="myModel", dataType="fp32"):
     for batch_idx, (data, target) in enumerate(train_loader):
         optimizer.zero_grad()
         if "bf16" == dataType:
-            with torch.cpu.amp.autocast():   # Auto Mixed Precision
+            with torch.amp.autocast('cpu'):   # Auto Mixed Precision
                 # Setting memory_format to torch.channels_last could improve performance with 4D input data. This is optional.
                 data = data.to(memory_format=torch.channels_last)
                 output = model(data)
@@ -106,7 +106,11 @@ def main():
 
     # Train models and acquire training times
     print("Training model with BF16 with AVX512")
-    bf16_noAmx_training_time = trainModel(train_loader, modelName="bf16_noAmx", dataType="bf16")
+    bf16_avx512_training_time = trainModel(train_loader, modelName="bf16_noAmx", dataType="bf16")
+
+    # Save variable
+    with open('bf16_noAmx_training_time.txt', 'w') as f:
+        f.write(str(bf16_avx512_training_time))
 
 if __name__ == '__main__':
     main()