MNIST plot first test figure after transform

user_15074450 0
Viewed 39

Here are some of my code, I want to know how to add code to plot the first figure from the test dataset after transforms?

transform=transforms.Compose([
        transforms.ToTensor(),
        AddGaussianNoise(0, 1),
        transforms.Normalize((0.1307,), (0.3081,))
        ])
    dataset1 = datasets.MNIST('../data', train=True, download=True,
                       transform=transform)
    dataset2 = datasets.MNIST('../data', train=False,
                       transform=transform)
    train_loader = torch.utils.data.DataLoader(dataset1,**train_kwargs)
    test_loader = torch.utils.data.DataLoader(dataset2, **test_kwargs)

    model = Net().to(device)
    optimizer = optim.Adadelta(model.parameters(), lr=args.lr)

    scheduler = StepLR(optimizer, step_size=1, gamma=args.gamma)
    for epoch in range(1, args.epochs + 1):
        train(args, model, device, train_loader, optimizer, epoch)
        test(model, device, test_loader)
        scheduler.step()

    if args.save_model:
        torch.save(model.state_dict(), "mnist_cnn.pt")
mnist pytorch deep-learning python-3.x python
1 Answers

Try this:

import torchvision.utils as vutils
import matplotlib.pyplot as plt

# get batch from dataloader -> (bs, ch, h, w)
imgs, lbls = next(iter(test_loader))

# make grid using images of the batch
img  = vutils.make_grid(imgs)

fig = plt.figure()
plt.axis("off")
# channels first to channels last conversion using permute
plt.imshow(img.permute(1, 2, 0))
plt.show()
user_141087340
Related