Documentation
Search…
PyTorch
PyTorch is one of the most popular frameworks for deep learning in Python, especially among researchers. W&B provides first class support for PyTorch, from logging gradients to profiling your code on the CPU and GPU.
Try our integration out in a colab notebook (with video walkthrough below) or see our example repo for scripts, including one on hyperparameter optimization using Hyperband on Fashion MNIST, plus the W&B Dashboard it generates.
Follow along with a video tutorial!
Logging gradients with wandb.watch
wandb.watch1
import wandb
2
wandb.init(config=args)
3
​
4
model = ... # set up your model
5
​
6
# Magic
7
wandb.watch(model, log_freq=100)
8
​
9
model.train()
10
for batch_idx, (data, target) in enumerate(train_loader):
11
output = model(data)
12
loss = F.nll_loss(output, target)
13
loss.backward()
14
optimizer.step()
15
if batch_idx % args.log_interval == 0:
16
wandb.log({"loss": loss})
Copied!
If you need to track multiple models in the same script, you can call
wandb.watch on each model separately. Reference documentation for this function is here.Gradients, metrics and the graph won't be logged until
wandb.log is called after a forward and backward pass.Logging images and media
You can pass PyTorch
Tensors with image data into wandb.Image and utilities from torchvision will be used to convert them to images automatically:1
images_t = ... # generate or load images as PyTorch Tensors
2
wandb.log({"examples" : [wandb.Image(im) for im in images_t]})
Copied!
For more on logging rich media to W&B in PyTorch and other frameworks, check out our media logging guide.
If you also want to include information alongside media, like your model's predictions or derived metrics, use a
wandb.Table.1
my_table = wandb.Table()
2
​
3
my_table.add_column("image", images_t)
4
my_table.add_column("label", labels)
5
my_table.add_column("class_prediction", predictions_t)
6
​
7
# Log your Table to W&B
8
wandb.log({"mnist_predictions": my_table})
Copied!
The code above generates a table like this one. This model's looking good!
Profiling PyTorch code
View detailed traces of PyTorch code execution inside W&B dashboards.
W&B integrates directly with PyTorch Kineto's Tensorboard plugin to provide tools for profiling PyTorch code, inspecting the details of CPU and GPU communication, and identifying bottlenecks and optimizations.
1
profile_dir = "path/to/run/tbprofile/"
2
profiler = torch.profiler.profile(
3
schedule=schedule, # see the profiler docs for details on scheduling
4
on_trace_ready=torch.profiler.tensorboard_trace_handler(profile_dir)
5
with_stack=True)
6
​
7
with profiler:
8
... # run the code you want to profile here
9
# see the profiler docs for detailed usage information
10
​
11
# create a wandb Artifact
12
profile_art = wandb.Artifact("trace", type="profile")
13
# add the pt.trace.json files to the Artifact
14
profile_art.add_file(glob.glob(profile_dir + ".pt.trace.json"))
15
# log the artifact
16
profile_art.save()
Copied!
The interactive trace viewing tool is based on the Chrome Trace Viewer, which works best with the Chrome browser.
Last modified 2mo ago