torch显示保存拼接数据集图片

先显示一个大的拼接的图片
再保存部分小图片

import torch
from torch import nn, optim
import torch.nn.functional as F
from torch.autograd import Variable
from torch.utils.data import DataLoader
from torchvision import transforms
from torchvision import datasets
import matplotlib.pyplot as plt
# from logger import Logger
import numpy as np
batch_size = 128

train_dataset = datasets.MNIST(
    root='data', train=True, transform=transforms.ToTensor(), download=True)

train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)

# 把图片拼接成一个大文件
for i, data in enumerate(train_loader, 1):
    img, label = data
    每一行的拼接图片 = img[0, 0, :, :].numpy()
    for count in range(1,10):
        oneImg = img[count, 0, :, :].numpy()
        每一行的拼接图片 = np.hstack((每一行的拼接图片, + oneImg))
    if i == 1:
        总拼接图片 =  每一行的拼接图片
    else:
        总拼接图片 = np.vstack((总拼接图片,   每一行的拼接图片))

    if i > 10:
        plt.imshow(总拼接图片)
        plt.imsave("temp/拼接图片.jpg", 总拼接图片, cmap="gray")
        plt.show()
        break

# 把图片挨个保存为文件
for i, data in enumerate(train_loader, 1):
    img, label = data
    for count in range(10):
        oneImg = img[count, 0, :, :]
        oneLabel = label.numpy()[count]
        plt.imsave("temp/" + str(oneLabel ) + "-" + str(i) + "-" +str(count) + ".jpg", oneImg, cmap="gray")
    if i > 5:  # 避免存太多文件
        break