Skip to content

Instantly share code, notes, and snippets.

@Tony607
Created May 9, 2019 13:37
Show Gist options
  • Save Tony607/ee50fec26e172310b128523001051c6b to your computer and use it in GitHub Desktop.
Save Tony607/ee50fec26e172310b128523001051c6b to your computer and use it in GitHub Desktop.
How to run Tensorboard for PyTorch 1.1.0 inside Jupyter notebook | DLology
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(False)
# 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)
images, labels = next(iter(trainloader))
grid = torchvision.utils.make_grid(images)
writer.add_image('images', grid, 0)
writer.add_graph(model, images)
writer.close()
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment