如何在TensorBoard中区分MPS/CPU设备的可视化结果？

要实现区分不同设备的TensorBoard记录，只需要做以下几个关键修改：

1. 循环内根据迭代的设备名称创建对应设备，而非用外部固定的device变量
2. 将设备名称拼接进add_image和add_graph的标签名中
3. 把writer.close()移到循环外，避免第一次循环就关闭写入器导致后续记录失败

修改后的完整代码：

import torch
import torchvision
from torch.utils.tensorboard import SummaryWriter
from torchvision import datasets, transforms

# Writer will output to ./runs/ directory by default
writer = SummaryWriter()

transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5,), (0.5,))])
trainset = datasets.MNIST('mnist_train', train=True, download=True, transform=transform)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=64, shuffle=True)
model = torchvision.models.resnet50(True)
# Have ResNet model take in grayscale rather than RGB
model.conv1 = torch.nn.Conv2d(1, 64, kernel_size=7, stride=2, padding=3, bias=False)

# 遍历目标设备列表
for device_name in ["mps","cpu"]:
    # 根据当前循环的设备名创建设备对象，同时判断mps可用性
    device = torch.device(device_name if (device_name == 'mps' and torch.has_mps) else 'cpu')
    model = model.to(device)
    images, labels = next(iter(trainloader))
    images = images.to(device)

    grid = torchvision.utils.make_grid(images)
    # 拼接设备名到图片标签
    writer.add_image(f'images_{device_name}', grid, 0)
    # 拼接设备名到图结构标签
    writer.add_graph(model, images, graph_name=f'model_{device_name}')

# 所有记录完成后再关闭写入器
writer.close()

关键修改说明： #

• 循环变量改为device_name，明确指代当前设备名称
• 每次循环根据device_name创建对应的device对象，同时做了mps可用性判断，避免mps不可用时报错
• 使用Python的f-string格式化标签名：f'images_{device_name}'和f'model_{device_name}'，这样TensorBoard里就会显示images_mps、images_cpu、model_mps、model_cpu的独立条目
• 将writer.close()移到循环外部，确保两次循环的记录都能被正确写入

内容的提问来源于stack exchange，提问作者massisenergy