AttributeError: module 'torch.cuda._memory_viz' has no attribute 'profile_plot'

1

I want to use torch.cuda._memory_viz.profile_plot to review my gpu usage. However, I got this error AttributeError: module 'torch.cuda._memory_viz' has no attribute 'profile_plot'.
My torch version is 1.13+cu117.
Anyone can help? Thanks.

2

profile_plot wasn’t available in torch==1.13.0 so you would need to update your PyTorch installation.

3

@ptrblck
Thanks. I have updated the version to 2.0.0 and find it. However, when I do
plot_html = profile_plot(prof, device=torch.device("cuda:0"))
it shows:

Traceback (most recent call last):
File “/media/lk2/sgccai/15_DLAcceleration/lukuan_lightning_test/torch_viz.py”, line 54, in
plot_html = profile_plot(prof, device=torch.device(“cuda:0”))
File “/home/lk2/anaconda3/envs/ai/lib/python3.10/site-packages/torch/cuda/_memory_viz.py”, line 415, in profile_plot
for event in memory_profile._op_tree.sorted_nodes:
AttributeError: ‘profile’ object has no attribute ‘_op_tree’

My whole code is:

import torch
import torch.nn
import torch.optim
import torch.profiler
import torch.utils.data
import torchvision.datasets
import torchvision.models
import torchvision.transforms as T
import os
from torch.cuda._memory_viz import profile_plot
import datetime

transform = T.Compose(
[T.Resize(224),
T.ToTensor(),
T.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])
train_set = torchvision.datasets.CIFAR10(root=‘./data’, train=True, download=True, transform=transform)
train_loader = torch.utils.data.DataLoader(train_set, batch_size=32, shuffle=True)

device = torch.device(“cuda:0”)
model = torchvision.models.resnet18(weights=‘IMAGENET1K_V1’).cuda(device)
criterion = torch.nn.CrossEntropyLoss().cuda(device)
optimizer = torch.optim.SGD(model.parameters(), lr=0.001, momentum=0.9)
model.train()

def train(data):
inputs, labels = data[0].to(device=device), data[1].to(device=device)
outputs = model(inputs)
loss = criterion(outputs, labels)
optimizer.zero_grad()
loss.backward()
optimizer.step()

with torch.profiler.profile(
schedule=torch.profiler.schedule(wait=1, warmup=1, active=3, repeat=2),
record_shapes=True,
profile_memory=True,
with_stack=True
) as prof:
for step, batch_data in enumerate(train_loader):
if step >= (1 + 1 + 3) * 2:
break
train(batch_data)
prof.step()

trace_dir = os.path.abspath(os.path.join(os.path.split(file)[0], “.”, “traces”))
now = datetime.datetime.now().strftime(“%Y_%m_%d:%H.%M.%S.pt.trace.json.gz”)
prof.export_chrome_trace(os.path.join(trace_dir, f"lukuan_{now}"))

plot_html = profile_plot(prof, device=torch.device(“cuda:0”))
with open(os.path.join(trace_dir, f"lukuan_{now}.html"), “w”) as f:
f.write(plot_html)