将迭代器中的值附加到对象中

Question

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)

device=torch.device('mps' if torch.has_mps else 'cpu')
for i in ["mps","cpu"]:
    model=model.to(device)
    images, labels = next(iter(trainloader))
    images=images.to(device)

    grid = torchvision.utils.make_grid(images.to(device))
    writer.add_image('images', grid, 0)
    writer.add_graph(model.to(device), images.to(device))
    writer.close()

我有这段代码，我想在

mps

和

cpu

对象中附加迭代器（

writer.add_image()

和

writer.add_graph()

）的值，这样我就可以区分哪个起源于哪个运行。换句话说，我想在

tensorboard

仪表板中看到

images_mps

＆images_cpu（通过对

writer.add_image('images', grid, 0)

行进行一些修改？）。

很抱歉这个菜鸟问题，但感谢任何帮助/建议！

Answer 1

Try to make 2

SummaryWriter

specifying 2

log_dir

, then use one writer for

mps

and other writer for `cpu
https://pytorch.org/docs/stable/tensorboard.html#torch.utils.tensorboard.writer.SummaryWriter