博客

import torch
import numpy as np

from torchvision import datasets
import torchvision.transforms as transforms

# Set the parameters
num_workers = 0
batch_size = 20

# Converting the Images to tensors using Transforms
transform = transforms.ToTensor()

train_data = datasets.MNIST(root='data', train=True,
                                   download=True, transform=transform)
test_data = datasets.MNIST(root='data', train=False,
                                  download=True, transform=transform)

# Loading the Data
train_loader = torch.utils.data.DataLoader(train_data, batch_size=batch_size,
    num_workers=num_workers)
test_loader = torch.utils.data.DataLoader(test_data, batch_size=batch_size,
    num_workers=num_workers)

#用matplotlib画图
import matplotlib.pyplot as plt
%matplotlib inline

#iter（）迭代器迭代batch_size个图像
dataiter = iter(train_loader)
images, labels = dataiter.next()
images = images.numpy()

# figsize指定figure的宽和高
fig = plt.figure(figsize=(25, 4))
for image in np.arange(20):
	# add_subplot(a,b,c) a为子图总行数，b为列数，c为画图的位置
    ax = fig.add_subplot(2, 20/2, image+1, xticks=[], yticks=[])
    # 从数组的形状中删除单维条目，即把shape中为1的维度去掉
    ax.imshow(np.squeeze(images[image]), cmap='gray')
    ax.set_title(str(labels[image].item()))