Skip to content

Commit bc6255b

Browse files
VoVAllenVoVAllen
authored
example: Add torch profiler example (tensorchord#1026)
* add profiler example Signed-off-by: Jinjing.Zhou <[email protected]> * lint Signed-off-by: Jinjing.Zhou <[email protected]> * fix Signed-off-by: Jinjing.Zhou <[email protected]> * fix markdown lint Signed-off-by: Jinjing.Zhou <[email protected]> Signed-off-by: Jinjing.Zhou <[email protected]>
1 parent ded3fce commit bc6255b

File tree

7 files changed

+190
-11
lines changed

7 files changed

+190
-11
lines changed

‎.github/workflows/link-check.yml

+12-11
Original file line numberDiff line numberDiff line change
@@ -5,21 +5,22 @@ on:
55
branches:
66
- main
77
paths:
8-
- '.github/workflows/**'
9-
- '**.md'
8+
- ".github/workflows/**"
9+
- "**.md"
1010
pull_request:
1111
paths:
12-
- '.github/workflows/**'
13-
- '**.md'
12+
- ".github/workflows/**"
13+
- "**.md"
1414

1515
jobs:
1616
markdown-link-check:
1717
runs-on: ubuntu-latest
1818
steps:
19-
- uses: actions/checkout@v3
20-
- uses: gaurav-nelson/github-action-markdown-link-check@v1
21-
with:
22-
file-path: 'README.md'
23-
folder-path: 'docs'
24-
check-modified-files-only: yes
25-
base-branch : main
19+
- uses: actions/checkout@v3
20+
- uses: gaurav-nelson/github-action-markdown-link-check@v1
21+
with:
22+
file-path: "README.md"
23+
folder-path: "docs"
24+
check-modified-files-only: yes
25+
base-branch: main
26+
config-file: .markdown-lint.json

‎.markdown-lint.json

+7
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,7 @@
1+
{
2+
"ignorePatterns": [
3+
{
4+
"pattern": "^http://localhost.*"
5+
},
6+
]
7+
}

‎examples/pytorch-profiler/README.md

+7
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,7 @@
1+
# PyTorch profiler example
2+
3+
This example is adopted from torch's [official tutorial](https://pytorch.org/tutorials/intermediate/tensorboard_profiler_tutorial.html). It shows how to use PyTorch profiler to analyze performance bottlenecks in a model.
4+
5+
## Usage
6+
7+
Run `envd up` at current folder. And execute `python main.py` in the envd container. Then you can see the profiling result through TensorBoard at http://localhost:8888

‎examples/pytorch-profiler/build.envd

+16
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,16 @@
1+
envdlib = include("https://github.com/tensorchord/envdlib")
2+
3+
4+
def build():
5+
base(os="ubuntu20.04", language="python3")
6+
shell("zsh")
7+
install.cuda(version="11.2.0", cudnn="8")
8+
install.python_packages(
9+
[
10+
"torch",
11+
"torchvision",
12+
"torch_tb_profiler",
13+
"--extra-index-url https://download.pytorch.org/whl/cu113",
14+
]
15+
)
16+
envdlib.tensorboard(envd_port=8888, envd_dir="/home/envd/log", host_port=8888)

‎examples/pytorch-profiler/main.py

+66
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,66 @@
1+
# Adopt from https://pytorch.org/tutorials/intermediate/tensorboard_profiler_tutorial.html
2+
import torch
3+
import torch.nn
4+
import torch.optim
5+
import torch.profiler
6+
import torch.utils.data
7+
import torchvision.datasets
8+
import torchvision.models
9+
import torchvision.transforms as T
10+
from torchvision.models import efficientnet_b0, EfficientNet_B0_Weights
11+
12+
transform = T.Compose(
13+
[T.Resize(224), T.ToTensor(), T.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))]
14+
)
15+
train_set = torchvision.datasets.CIFAR10(
16+
root="./data", train=True, download=True, transform=transform
17+
)
18+
train_loader = torch.utils.data.DataLoader(train_set, batch_size=2, shuffle=True)
19+
20+
device = torch.device("cuda:0")
21+
model = efficientnet_b0(weights=EfficientNet_B0_Weights.DEFAULT).to(device)
22+
criterion = torch.nn.CrossEntropyLoss().cuda(device)
23+
optimizer = torch.optim.SGD(model.parameters(), lr=0.001, momentum=0.9)
24+
model.train()
25+
26+
27+
def train(data):
28+
inputs, labels = data[0].to(device=device), data[1].to(device=device)
29+
outputs = model(inputs)
30+
loss = criterion(outputs, labels)
31+
optimizer.zero_grad()
32+
loss.backward()
33+
optimizer.step()
34+
35+
36+
with torch.profiler.profile(
37+
schedule=torch.profiler.schedule(wait=1, warmup=1, active=3, repeat=2),
38+
on_trace_ready=torch.profiler.tensorboard_trace_handler(
39+
"/home/envd/log/efficientnet"
40+
),
41+
record_shapes=True,
42+
profile_memory=True,
43+
with_stack=True,
44+
) as prof:
45+
for step, batch_data in enumerate(train_loader):
46+
if step >= (1 + 1 + 3) * 2:
47+
break
48+
train(batch_data)
49+
prof.step() # Need to call this at the end of each step to notify profiler of steps' boundary.
50+
51+
52+
prof = torch.profiler.profile(
53+
schedule=torch.profiler.schedule(wait=1, warmup=1, active=3, repeat=2),
54+
on_trace_ready=torch.profiler.tensorboard_trace_handler(
55+
"/home/envd/log/efficientnet"
56+
),
57+
record_shapes=True,
58+
with_stack=True,
59+
)
60+
prof.start()
61+
for step, batch_data in enumerate(train_loader):
62+
if step >= (1 + 1 + 3) * 2:
63+
break
64+
train(batch_data)
65+
prof.step()
66+
prof.stop()

‎examples/pytorch_profiler/build.envd

+16
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,16 @@
1+
envdlib = include("https://github.com/tensorchord/envdlib")
2+
3+
4+
def build():
5+
base(os="ubuntu20.04", language="python3")
6+
shell("zsh")
7+
install.cuda(version="11.2.0", cudnn="8")
8+
install.python_packages(
9+
[
10+
"torch",
11+
"torchvision",
12+
"torch_tb_profiler",
13+
"--extra-index-url https://download.pytorch.org/whl/cu113",
14+
]
15+
)
16+
envdlib.tensorboard(8888, envd_dir="/home/envd/log/efficientnet")

‎examples/pytorch_profiler/main.py

+66
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,66 @@
1+
# Adopt from https://pytorch.org/tutorials/intermediate/tensorboard_profiler_tutorial.html
2+
import torch
3+
import torch.nn
4+
import torch.optim
5+
import torch.profiler
6+
import torch.utils.data
7+
import torchvision.datasets
8+
import torchvision.models
9+
import torchvision.transforms as T
10+
from torchvision.models import efficientnet_b0, EfficientNet_B0_Weights
11+
12+
transform = T.Compose(
13+
[T.Resize(224), T.ToTensor(), T.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))]
14+
)
15+
train_set = torchvision.datasets.CIFAR10(
16+
root="./data", train=True, download=True, transform=transform
17+
)
18+
train_loader = torch.utils.data.DataLoader(train_set, batch_size=2, shuffle=True)
19+
20+
device = torch.device("cuda:0")
21+
model = efficientnet_b0(weights=EfficientNet_B0_Weights.DEFAULT).to(device)
22+
criterion = torch.nn.CrossEntropyLoss().cuda(device)
23+
optimizer = torch.optim.SGD(model.parameters(), lr=0.001, momentum=0.9)
24+
model.train()
25+
26+
27+
def train(data):
28+
inputs, labels = data[0].to(device=device), data[1].to(device=device)
29+
outputs = model(inputs)
30+
loss = criterion(outputs, labels)
31+
optimizer.zero_grad()
32+
loss.backward()
33+
optimizer.step()
34+
35+
36+
with torch.profiler.profile(
37+
schedule=torch.profiler.schedule(wait=1, warmup=1, active=3, repeat=2),
38+
on_trace_ready=torch.profiler.tensorboard_trace_handler(
39+
"/home/envd/log/efficientnet"
40+
),
41+
record_shapes=True,
42+
profile_memory=True,
43+
with_stack=True,
44+
) as prof:
45+
for step, batch_data in enumerate(train_loader):
46+
if step >= (1 + 1 + 3) * 2:
47+
break
48+
train(batch_data)
49+
prof.step() # Need to call this at the end of each step to notify profiler of steps' boundary.
50+
51+
52+
prof = torch.profiler.profile(
53+
schedule=torch.profiler.schedule(wait=1, warmup=1, active=3, repeat=2),
54+
on_trace_ready=torch.profiler.tensorboard_trace_handler(
55+
"/home/envd/log/efficientnet"
56+
),
57+
record_shapes=True,
58+
with_stack=True,
59+
)
60+
prof.start()
61+
for step, batch_data in enumerate(train_loader):
62+
if step >= (1 + 1 + 3) * 2:
63+
break
64+
train(batch_data)
65+
prof.step()
66+
prof.stop()

0 commit comments

Comments
 (0)