utils.py

import torch
import matplotlib.pyplot as plt
import numpy as np
from torchvision.utils import make_grid

# device = 'cuda' if torch.cuda.is_available() else 'cpu'
device = 'cpu'
plt.interactive(False)


def show(img):
    npimg = img.numpy()
    plt.imshow(np.transpose(npimg, (1, 2, 0)), interpolation='nearest')
    plt.show()
    

def show_dataset(dataset, n=6):
    imgs = [dataset[i][0] for i in range(n)]
    grid = make_grid(imgs)
    plt.imshow(grid.numpy().transpose((1, 2, 0)))
    plt.tight_layout()
    plt.show()


def show_dl(dl, n=6):
    batch = None
    for batch in dl:
        break
    imgs = batch[0][:n]
    grid = make_grid(imgs)
    plt.imshow(grid.numpy().transpose((1, 2, 0)))
    plt.tight_layout()
    plt.show()