From 5fcf512011620993b2e4c3625d1dcda2c0e19df3 Mon Sep 17 00:00:00 2001
From: "Tsai, Louie" <louie.tsai@intel.com>
Date: Tue, 22 Oct 2024 16:57:47 -0700
Subject: [PATCH] fixes some runtime issues

---
 .../Profiler_and_Optimization.ipynb           | 47 +++++++++++++++----
 1 file changed, 38 insertions(+), 9 deletions(-)

diff --git a/PyTorch/Profiling_and_Optimization/Profiler_and_Optimization.ipynb b/PyTorch/Profiling_and_Optimization/Profiler_and_Optimization.ipynb
index ab2eb66..546e463 100644
--- a/PyTorch/Profiling_and_Optimization/Profiler_and_Optimization.ipynb
+++ b/PyTorch/Profiling_and_Optimization/Profiler_and_Optimization.ipynb
@@ -378,6 +378,17 @@
     "The collected trace files will be saved to ./hpu_profile"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "c97a9142-a80e-4c4f-82b7-4f08509996c5",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from optimum.habana.utils import HabanaProfile\n",
+    "HabanaProfile.enable()"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": 5,
@@ -681,7 +692,7 @@
     }
    ],
    "source": [
-    "!python run_image_classification.py \\\n",
+    "%run run_image_classification.py \\\n",
     "    --model_name_or_path microsoft/swin-base-patch4-window7-224-in22k \\\n",
     "    --dataset_name cifar10 \\\n",
     "    --output_dir /tmp/outputs/ \\\n",
@@ -710,6 +721,16 @@
     "    --profiling_steps 3"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "4f029263-187b-46be-b017-09c5be1ef8dc",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "!mv ./hpu_profile/*.pt.trace.json ./hpu_profile/UNOPT.pt.trace.json"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": 7,
@@ -738,8 +759,7 @@
     }
    ],
    "source": [
-    "%cd hpu_profile\n",
-    "%ls -al"
+    "%ls -al hpu_profile"
    ]
   },
   {
@@ -790,7 +810,7 @@
    ],
    "source": [
     "%load_ext tensorboard\n",
-    "%tensorboard --logdir=~/Gaudi-tutorials/PyTorch/Profiling_and_Optimization/optimum-habana/examples/image-classification/hpu_profile --port 6006    # Your port selection may vary, default is 6006"
+    "%tensorboard --logdir=./hpu_profile --port 6006    # Your port selection may vary, default is 6006"
    ]
   },
   {
@@ -834,7 +854,7 @@
     }
    ],
    "source": [
-    "!habana_perf_tool --trace /root/Gaudi-tutorials/PyTorch/Profiling_and_Optimization/optimum-habana/examples/image-classification/hpu_profile/hls2-srv01-demolab_1924.1710802005850814532.pt.trace.json"
+    "!habana_perf_tool --trace ./hpu_profile/UNOPT.pt.trace.json"
    ]
   },
   {
@@ -857,9 +877,18 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "%cd ..\n",
-    "!hl-prof-config -e off -phase=multi-enq -g 1-20 -s my_profiling_session\n",
-    "!export HABANA_PROFILE=1"
+    "!hl-prof-config -e off -phase=multi-enq -g 1-20 -s my_profiling_session"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "1eb9fe7f-5817-4803-90b8-244da14213c0",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "os.environ[\"HABANA_PROFILE\"] = \"1\""
    ]
   },
   {
@@ -869,7 +898,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "!HABANA_PROFILE=1 python run_image_classification.py \\\n",
+    "%run run_image_classification.py \\\n",
     "    --model_name_or_path microsoft/swin-base-patch4-window7-224-in22k \\\n",
     "    --dataset_name cifar10 \\\n",
     "    --output_dir /tmp/outputs/ \\\n",