pytorch加载tiny ImageNet

# 使用foolbox实现JSMA算法
import foolbox
import torch
import numpy as np
from torchvision import datasets,models,transforms
import os
import matplotlib.pyplot as plt

device = torch.device("cuda" if (torch.cuda.is_available()) else "cpu")

# instantiate model
kmodel = models.resnet18(pretrained=True).eval().to(device)

def tinyImageNet_sample(batch_size, data_dir, model):
    num_label = 200
    # normalize = transforms.Normalize((0.4802, 0.4481, 0.3975), (0.2770, 0.2691, 0.2821))

    # 将图像归一化到 -1~1，不写则归一化到 0~1
    normalize = transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))

    transform_train = transforms.Compose([transforms.Resize(224), transforms.ToTensor(), ])

    # transform_test = transforms.Compose([transforms.Resize(32), transforms.ToTensor(), normalize, ])

    trainset = datasets.ImageFolder(root=os.path.join(data_dir, 'train'), transform=transform_train)

    # testset = datasets.ImageFolder(root=os.path.join(data_dir, 'val'), transform=transform_test)

    train_loader = torch.utils.data.DataLoader(trainset, batch_size=batch_size, shuffle=True, pin_memory=True)
    # test_loader = torch.utils.data.DataLoader(testset, batch_size=batch_size, shuffle=False, pin_memory=True)
    # return train_loader, test_loader, num_label

    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    images = None
    labels = None

    for data in train_loader:
        images,labels = data
        break

    images = images.to(device)
    with torch.no_grad():
        _,predict = torch.max(model(images),1)
    return images.data.cpu(), predict.data.cpu()

images, labels = tinyImageNet_sample(10, './Data/tiny-imagenet-200',kmodel)

def plot_clean(n_cols,n_rows,clean_example,clean_target,attack_name):
    n_cols = n_cols
    n_rows = n_rows
    cnt = 1
    cnt1 = 1
    plt.figure(figsize=(n_rows,n_cols))
    for i in range(n_cols):
        for j in range(n_rows):
            plt.subplot(n_cols,n_rows,cnt1)
            plt.xticks([])
            plt.yticks([])
            if j == 0:
                plt.ylabel(attack_name,size=15)
            plt.title("{}".format(clean_target[cnt-1]))
            plt.imshow(clean_example[cnt-1].to('cpu').detach().numpy().transpose(1,2,0))
            plt.subplot(n_cols,n_rows,cnt1)
            cnt1 = cnt1 + 1
            cnt = cnt + 1
    plt.show()
    print('n')

plot_clean(2,5,images,labels,'CW targeted')