Skip to content

Commit 1ef0d11

Browse files
authored
Add non-context manager tutorial for PyTorch Tensorboard plugin (#1723)
* Add non-context manager tutorial for PyTorch Tensorboard plugin * add split line in tutorials * fix bugs
1 parent 2c599be commit 1ef0d11

File tree

1 file changed

+14
-0
lines changed

1 file changed

+14
-0
lines changed

intermediate_source/tensorboard_profiler_tutorial.py

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -132,6 +132,20 @@ def train(data):
132132
train(batch_data)
133133
prof.step() # Need to call this at the end of each step to notify profiler of steps' boundary.
134134

135+
######################################################################
136+
# Alternatively, the following non-context manager start/stop is supported as well.
137+
prof = torch.profiler.profile(
138+
schedule=torch.profiler.schedule(wait=1, warmup=1, active=3, repeat=2),
139+
on_trace_ready=torch.profiler.tensorboard_trace_handler('./log/resnet18'),
140+
record_shapes=True,
141+
with_stack=True)
142+
prof.start()
143+
for step, batch_data in enumerate(train_loader):
144+
if step >= (1 + 1 + 3) * 2:
145+
break
146+
train(batch_data)
147+
prof.step()
148+
prof.stop()
135149

136150
######################################################################
137151
# 3. Run the profiler

0 commit comments

Comments
 (0)