Files
mrdbourke--pytorch-deep-lea…/docs/08_pytorch_profiling.ipynb
T
wehub-resource-sync 37c25cd088
make_docs / deploy (push) Waiting to run
chore: import upstream snapshot with attribution
2026-07-13 13:27:58 +08:00

692 lines
21 KiB
Plaintext

{
"cells": [
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# 08: PyTorch Profiling\n",
"\n",
"This notebook is an experiment to try out the PyTorch profiler.\n",
"\n",
"See here for more:\n",
"* https://pytorch.org/blog/introducing-pytorch-profiler-the-new-and-improved-performance-tool/\n",
"* https://pytorch.org/docs/stable/profiler.html"
]
},
{
"cell_type": "code",
"execution_count": 10,
"metadata": {},
"outputs": [],
"source": [
"import torch\n",
"import torchvision\n",
"from torch import nn\n",
"from torchvision import transforms, datasets\n",
"from torchinfo import summary\n",
"\n",
"import numpy as np\n",
"import matplotlib.pyplot as plt\n",
"\n",
"from going_modular import data_setup, engine"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Setup device"
]
},
{
"cell_type": "code",
"execution_count": 11,
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"'cuda'"
]
},
"execution_count": 11,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
"device"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Get and load data"
]
},
{
"cell_type": "code",
"execution_count": 12,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"data/10_whole_foods dir exists, skipping download\n"
]
}
],
"source": [
"import os\n",
"import requests\n",
"from zipfile import ZipFile\n",
"\n",
"def get_food_image_data():\n",
" if not os.path.exists(\"data/10_whole_foods\"):\n",
" os.makedirs(\"data/\", exist_ok=True)\n",
" # Download data\n",
" data_url = \"https://storage.googleapis.com/food-vision-image-playground/10_whole_foods.zip\"\n",
" print(f\"Downloading data from {data_url}...\")\n",
" requests.get(data_url)\n",
" # Unzip data\n",
" targ_dir = \"data/10_whole_foods\"\n",
" print(f\"Extracting data to {targ_dir}...\")\n",
" with ZipFile(\"10_whole_foods.zip\") as zip_ref:\n",
" zip_ref.extractall(targ_dir)\n",
" else:\n",
" print(\"data/10_whole_foods dir exists, skipping download\")\n",
"\n",
"get_food_image_data()"
]
},
{
"cell_type": "code",
"execution_count": 38,
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"(<torch.utils.data.dataloader.DataLoader at 0x7f052ab26e20>,\n",
" <torch.utils.data.dataloader.DataLoader at 0x7f05299eed00>,\n",
" ['apple',\n",
" 'banana',\n",
" 'beef',\n",
" 'blueberries',\n",
" 'carrots',\n",
" 'chicken_wings',\n",
" 'egg',\n",
" 'honey',\n",
" 'mushrooms',\n",
" 'strawberries'])"
]
},
"execution_count": 38,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"# Setup dirs\n",
"train_dir = \"data/10_whole_foods/train\"\n",
"test_dir = \"data/10_whole_foods/test\"\n",
"\n",
"# Setup ImageNet normalization levels (turns all images into similar distribution as ImageNet)\n",
"normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],\n",
" std=[0.229, 0.224, 0.225])\n",
"\n",
"# Create starter transform\n",
"simple_transform = transforms.Compose([\n",
" transforms.Resize((224, 224)),\n",
" transforms.ToTensor(),\n",
" normalize\n",
"]) \n",
"\n",
"# Create data loaders\n",
"train_dataloader, test_dataloader, class_names = data_setup.create_dataloaders(\n",
" train_dir=train_dir,\n",
" test_dir=test_dir,\n",
" transform=simple_transform,\n",
" batch_size=32,\n",
" num_workers=8\n",
")\n",
"\n",
"train_dataloader, test_dataloader, class_names"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Load model "
]
},
{
"cell_type": "code",
"execution_count": 66,
"metadata": {},
"outputs": [],
"source": [
"model = torchvision.models.efficientnet_b0(pretrained=True).to(device)\n",
"# model"
]
},
{
"cell_type": "code",
"execution_count": 67,
"metadata": {},
"outputs": [],
"source": [
"# Update the classifier\n",
"model.classifier = torch.nn.Sequential(\n",
" nn.Dropout(p=0.2),\n",
" nn.Linear(1280, len(class_names)).to(device))\n",
"\n",
"# Freeze all base layers \n",
"for param in model.features.parameters():\n",
" param.requires_grad = False"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Train model and track results"
]
},
{
"cell_type": "code",
"execution_count": 68,
"metadata": {},
"outputs": [],
"source": [
"# Define loss and optimizer\n",
"loss_fn = nn.CrossEntropyLoss()\n",
"optimizer = torch.optim.Adam(model.parameters(), lr=0.001)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### Adjust training function to track results with `SummaryWriter`"
]
},
{
"cell_type": "code",
"execution_count": 69,
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"'EfficietNetB0'"
]
},
"execution_count": 69,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"model.name = \"EfficietNetB0\"\n",
"model.name"
]
},
{
"cell_type": "code",
"execution_count": 70,
"metadata": {},
"outputs": [],
"source": [
"from torch.utils.tensorboard import SummaryWriter\n",
"from going_modular.engine import train_step, test_step\n",
"from tqdm import tqdm\n",
"writer = SummaryWriter()"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Update the `train_step()` function to include the PyTorch profiler."
]
},
{
"cell_type": "code",
"execution_count": 71,
"metadata": {},
"outputs": [],
"source": [
"def train_step(model, dataloader, loss_fn, optimizer):\n",
" model.train()\n",
" train_loss, train_acc = 0, 0\n",
" ## NEW: Add PyTorch profiler\n",
"\n",
" dir_to_save_logs = os.path.join(\"logs\", datetime.now().strftime(\"%Y-%m-%d-%H-%M\"))\n",
" with torch.profiler.profile(\n",
" on_trace_ready=torch.profiler.tensorboard_trace_handler(dir_name=dir_to_save_logs),\n",
" # with_stack=True # this adds a lot of overhead to training (tracing all the stack)\n",
" ):\n",
" for batch, (X, y) in enumerate(dataloader):\n",
" # Send data to GPU\n",
" X, y = X.to(device, non_blocking=True), y.to(device, non_blocking=True)\n",
" \n",
" # Turn on mixed precision if available\n",
" with torch.autocast(device_type=device, enabled=True):\n",
" # 1. Forward pass\n",
" y_pred = model(X)\n",
"\n",
" # 2. Calculate loss\n",
" loss = loss_fn(y_pred, y)\n",
"\n",
" # 3. Optimizer zero grad\n",
" optimizer.zero_grad()\n",
"\n",
" # 4. Loss backward\n",
" loss.backward()\n",
"\n",
" # 5. Optimizer step\n",
" optimizer.step()\n",
"\n",
" # 6. Calculate metrics\n",
" train_loss += loss.item()\n",
" y_pred_class = torch.softmax(y_pred, dim=1).argmax(dim=1)\n",
" # print(f\"y: \\n{y}\\ny_pred_class:{y_pred_class}\")\n",
" # print(f\"y argmax: {y_pred.argmax(dim=1)}\")\n",
" # print(f\"Equal: {(y_pred_class == y)}\")\n",
" train_acc += (y_pred_class == y).sum().item() / len(y_pred)\n",
" # print(f\"batch: {batch} train_acc: {train_acc}\")\n",
"\n",
" # Adjust returned metrics\n",
" return train_loss / len(dataloader), train_acc / len(dataloader)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"TK - Now to use the writer, we've got to adjust the `train()` function..."
]
},
{
"cell_type": "code",
"execution_count": 72,
"metadata": {},
"outputs": [],
"source": [
"def train(\n",
" model,\n",
" train_dataloader,\n",
" test_dataloader,\n",
" optimizer,\n",
" loss_fn=nn.CrossEntropyLoss(),\n",
" epochs=5,\n",
"):\n",
"\n",
" results = {\"train_loss\": [], \"train_acc\": [], \"test_loss\": [], \"test_acc\": []}\n",
"\n",
" for epoch in tqdm(range(epochs)):\n",
" train_loss, train_acc = train_step(\n",
" model=model,\n",
" dataloader=train_dataloader,\n",
" loss_fn=loss_fn,\n",
" optimizer=optimizer,\n",
" )\n",
" test_loss, test_acc = test_step(\n",
" model=model, dataloader=test_dataloader, loss_fn=loss_fn\n",
" )\n",
"\n",
" # Print out what's happening\n",
" print(\n",
" f\"Epoch: {epoch+1} | \"\n",
" f\"train_loss: {train_loss:.4f} | \"\n",
" f\"train_acc: {train_acc:.4f} | \"\n",
" f\"test_loss: {test_loss:.4f} | \"\n",
" f\"test_acc: {test_acc:.4f}\"\n",
" )\n",
"\n",
" # Update results\n",
" results[\"train_loss\"].append(train_loss)\n",
" results[\"train_acc\"].append(train_acc)\n",
" results[\"test_loss\"].append(test_loss)\n",
" results[\"test_acc\"].append(test_acc)\n",
"\n",
" # Add results to SummaryWriter\n",
" writer.add_scalars(main_tag=\"Loss\", \n",
" tag_scalar_dict={\"train_loss\": train_loss,\n",
" \"test_loss\": test_loss},\n",
" global_step=epoch)\n",
" writer.add_scalars(main_tag=\"Accuracy\", \n",
" tag_scalar_dict={\"train_acc\": train_acc,\n",
" \"test_acc\": test_acc}, \n",
" global_step=epoch)\n",
" \n",
" # Close the writer\n",
" writer.close()\n",
"\n",
" return results"
]
},
{
"cell_type": "code",
"execution_count": 73,
"metadata": {},
"outputs": [
{
"name": "stderr",
"output_type": "stream",
"text": [
" 20%|██ | 1/5 [00:05<00:21, 5.27s/it]"
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"Epoch: 1 | train_loss: 1.9644 | train_acc: 0.4386 | test_loss: 1.5205 | test_acc: 0.7865\n"
]
},
{
"name": "stderr",
"output_type": "stream",
"text": [
" 40%|████ | 2/5 [00:09<00:14, 4.94s/it]"
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"Epoch: 2 | train_loss: 1.2589 | train_acc: 0.7878 | test_loss: 1.1589 | test_acc: 0.7604\n"
]
},
{
"name": "stderr",
"output_type": "stream",
"text": [
" 60%|██████ | 3/5 [00:14<00:09, 4.72s/it]"
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"Epoch: 3 | train_loss: 0.8642 | train_acc: 0.8776 | test_loss: 0.9347 | test_acc: 0.7917\n"
]
},
{
"name": "stderr",
"output_type": "stream",
"text": [
" 80%|████████ | 4/5 [00:18<00:04, 4.56s/it]"
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"Epoch: 4 | train_loss: 0.6827 | train_acc: 0.8856 | test_loss: 0.6637 | test_acc: 0.8750\n"
]
},
{
"name": "stderr",
"output_type": "stream",
"text": [
"100%|██████████| 5/5 [00:23<00:00, 4.65s/it]"
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"Epoch: 5 | train_loss: 0.5688 | train_acc: 0.9069 | test_loss: 0.6175 | test_acc: 0.8854\n"
]
},
{
"name": "stderr",
"output_type": "stream",
"text": [
"\n"
]
}
],
"source": [
"# Train model\n",
"# Note: Not using engine.train() since the original script isn't updated\n",
"results = train(model=model,\n",
" train_dataloader=train_dataloader,\n",
" test_dataloader=test_dataloader,\n",
" optimizer=optimizer,\n",
" loss_fn=loss_fn,\n",
" epochs=5)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Looks like mixed precision doesn't offer much benefit for smaller feature extraction models..."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# # Without mixed precision\n",
"# 20%|██ | 1/5 [00:03<00:14, 3.71s/it]Epoch: 1 | train_loss: 0.5229 | train_acc: 0.9054 | test_loss: 0.5776 | test_acc: 0.8542\n",
"# 40%|████ | 2/5 [00:07<00:11, 3.74s/it]Epoch: 2 | train_loss: 0.4699 | train_acc: 0.9001 | test_loss: 0.5160 | test_acc: 0.8802\n",
"# 60%|██████ | 3/5 [00:10<00:07, 3.63s/it]Epoch: 3 | train_loss: 0.3913 | train_acc: 0.9196 | test_loss: 0.4888 | test_acc: 0.8906\n",
"# 80%|████████ | 4/5 [00:14<00:03, 3.61s/it]Epoch: 4 | train_loss: 0.3724 | train_acc: 0.9371 | test_loss: 0.4931 | test_acc: 0.8698\n",
"# 100%|██████████| 5/5 [00:18<00:00, 3.61s/it]Epoch: 5 | train_loss: 0.3315 | train_acc: 0.9381 | test_loss: 0.4405 | test_acc: 0.8750\n",
"\n",
"# # With mixed precision\n",
"# 20%|██ | 1/5 [00:04<00:17, 4.40s/it]Epoch: 1 | train_loss: 0.3027 | train_acc: 0.9554 | test_loss: 0.4386 | test_acc: 0.8802\n",
"# 40%|████ | 2/5 [00:08<00:13, 4.49s/it]Epoch: 2 | train_loss: 0.2826 | train_acc: 0.9539 | test_loss: 0.4080 | test_acc: 0.8802\n",
"# 60%|██████ | 3/5 [00:13<00:08, 4.48s/it]Epoch: 3 | train_loss: 0.2450 | train_acc: 0.9609 | test_loss: 0.4130 | test_acc: 0.8750\n",
"# 80%|████████ | 4/5 [00:18<00:04, 4.53s/it]Epoch: 4 | train_loss: 0.2450 | train_acc: 0.9594 | test_loss: 0.4158 | test_acc: 0.8802\n",
"# 100%|██████████| 5/5 [00:22<00:00, 4.49s/it]Epoch: 5 | train_loss: 0.2307 | train_acc: 0.9639 | test_loss: 0.4124 | test_acc: 0.8906"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Try mixed precision with larger model\n",
"\n",
"Now we'll try turn on mixed precision with a larger model (e.g. EffifientNetB0 with all layers tuneable)."
]
},
{
"cell_type": "code",
"execution_count": 74,
"metadata": {},
"outputs": [],
"source": [
"# Unfreeze all base layers \n",
"for param in model.features.parameters():\n",
" param.requires_grad = True\n",
"\n",
"# for param in model.features.parameters():\n",
"# print(param.requires_grad)"
]
},
{
"cell_type": "code",
"execution_count": 75,
"metadata": {},
"outputs": [
{
"name": "stderr",
"output_type": "stream",
"text": [
" 20%|██ | 1/5 [00:13<00:53, 13.27s/it]"
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"Epoch: 1 | train_loss: 0.4934 | train_acc: 0.8586 | test_loss: 0.6467 | test_acc: 0.7969\n"
]
},
{
"name": "stderr",
"output_type": "stream",
"text": [
" 40%|████ | 2/5 [00:27<00:42, 14.09s/it]"
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"Epoch: 2 | train_loss: 0.1750 | train_acc: 0.9628 | test_loss: 1.1806 | test_acc: 0.8385\n"
]
},
{
"name": "stderr",
"output_type": "stream",
"text": [
" 60%|██████ | 3/5 [00:42<00:28, 14.28s/it]"
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"Epoch: 3 | train_loss: 0.1362 | train_acc: 0.9619 | test_loss: 0.5831 | test_acc: 0.8802\n"
]
},
{
"name": "stderr",
"output_type": "stream",
"text": [
" 80%|████████ | 4/5 [00:57<00:14, 14.47s/it]"
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"Epoch: 4 | train_loss: 0.1743 | train_acc: 0.9462 | test_loss: 0.5702 | test_acc: 0.8854\n"
]
},
{
"name": "stderr",
"output_type": "stream",
"text": [
"100%|██████████| 5/5 [01:11<00:00, 14.38s/it]"
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"Epoch: 5 | train_loss: 0.2437 | train_acc: 0.9352 | test_loss: 0.7096 | test_acc: 0.8125\n"
]
},
{
"name": "stderr",
"output_type": "stream",
"text": [
"\n"
]
}
],
"source": [
"# Train model\n",
"# Note: Not using engine.train() since the original script isn't updated\n",
"results = train(model=model,\n",
" train_dataloader=train_dataloader,\n",
" test_dataloader=test_dataloader,\n",
" optimizer=optimizer,\n",
" loss_fn=loss_fn,\n",
" epochs=5)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# # Without mixed precision...\n",
"# 20%|██ | 1/5 [00:11<00:46, 11.61s/it]Epoch: 1 | train_loss: 0.4507 | train_acc: 0.8648 | test_loss: 1.0603 | test_acc: 0.7604\n",
"# 40%|████ | 2/5 [00:24<00:36, 12.21s/it]Epoch: 2 | train_loss: 0.1659 | train_acc: 0.9464 | test_loss: 0.6398 | test_acc: 0.8490\n",
"# 60%|██████ | 3/5 [00:36<00:24, 12.38s/it]Epoch: 3 | train_loss: 0.1261 | train_acc: 0.9698 | test_loss: 0.7149 | test_acc: 0.8542\n",
"# 80%|████████ | 4/5 [00:49<00:12, 12.53s/it]Epoch: 4 | train_loss: 0.1250 | train_acc: 0.9609 | test_loss: 0.7441 | test_acc: 0.7917\n",
"# 100%|██████████| 5/5 [01:02<00:00, 12.42s/it]Epoch: 5 | train_loss: 0.1282 | train_acc: 0.9564 | test_loss: 0.8701 | test_acc: 0.8385\n",
"\n",
"# # With mixed precision...\n",
"# 20%|██ | 1/5 [00:13<00:53, 13.27s/it]Epoch: 1 | train_loss: 0.4934 | train_acc: 0.8586 | test_loss: 0.6467 | test_acc: 0.7969\n",
"# 40%|████ | 2/5 [00:27<00:42, 14.09s/it]Epoch: 2 | train_loss: 0.1750 | train_acc: 0.9628 | test_loss: 1.1806 | test_acc: 0.8385\n",
"# 60%|██████ | 3/5 [00:42<00:28, 14.28s/it]Epoch: 3 | train_loss: 0.1362 | train_acc: 0.9619 | test_loss: 0.5831 | test_acc: 0.8802\n",
"# 80%|████████ | 4/5 [00:57<00:14, 14.47s/it]Epoch: 4 | train_loss: 0.1743 | train_acc: 0.9462 | test_loss: 0.5702 | test_acc: 0.8854\n",
"# 100%|██████████| 5/5 [01:11<00:00, 14.38s/it]Epoch: 5 | train_loss: 0.2437 | train_acc: 0.9352 | test_loss: 0.7096 | test_acc: 0.8125"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Checking the PyTorch profiler, it seems that mixed precision utilises some Tensor Cores, however, these aren't large numbers.\n",
"\n",
"E.g. it uses 9-12% Tensor Cores. Perhaps the slow down when using mixed precision is because the tensors have to get altered and converted when there isn't very many of them. For example only 9-12% of tensors get converted so the speed up gains aren't realised on these tensors because they get cancelled out by the conversion time."
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": []
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Extensions\n",
"* Does changing the data input size to EfficientNetB4 change its results? E.g. input image size of (380, 380) instead of (224, 224)?"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
" "
]
}
],
"metadata": {
"interpreter": {
"hash": "3fbe1355223f7b2ffc113ba3ade6a2b520cadace5d5ec3e828c83ce02eb221bf"
},
"kernelspec": {
"display_name": "Python 3.9.7 64-bit (conda)",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.9.7"
},
"orig_nbformat": 4
},
"nbformat": 4,
"nbformat_minor": 2
}