blog_show-plant-dataloader.py

fig = plt.figure(figsize=(3, 7))
for imgs, lbs in dm.val_dataloader():
    # some stats about the batch - label distribution
    print(f'batch labels: {torch.sum(lbs, axis=0)}')
    print(f'image size: {imgs[0].shape}')
    # similar as above show just first images from the batch
    for i in range(3):
        ax = fig.add_subplot(3, 1, i + 1, xticks=[], yticks=[])
        ax.imshow(np.rollaxis(imgs[i].numpy(), 0, 3))
        ax.set_title(lbs[i])
    # stop after the first batch
    break