PyTorch 使用自己的数据集

  1. from __future__ import print_function
  3. import torch
  4. import torch.nn as nn
  5. import torch.nn.functional as F
  6. import torch.autograd.variable as Variable
  7. import torch.optim as optim
  9. import torchvision
  10. import torchvision.transforms as transforms
  12. import numpy as np
  13. import matplotlib.pyplot as plt
  1. %matplotlib inline

保存 CIFAR10 为图片

CIFAR10 数据集说明

  • 包含有 data_batch_1,…,data_batch_5,test_batch
  • 每一个batch文件包括一个字典，字典的元素是：
    • data：一个尺寸为 10000×3072，数据格式为 uint8 的 numpy array，每一行数据存储了一张 32×32 彩色图片的数据，前 1024 位是图像的红色通道数据，接着是绿色通道和蓝色通道。
    • label：一个包含 10000 个 0-9 数字的列表，对应 data 里每张图片的标签。

读取数据集并保存

重复执行会清空csv文件

  1. from matplotlib.image import imsave
  2. import numpy as np
  3. import csv
  4. import os
  7. def unpickle(file):
  8.     import pickle
  9.     with open(file, 'rb') as fo:
  10.         dict = pickle.load(fo, encoding='bytes')
  11.     return dict
  14. path = './data/cifar10/cifar-10-batches-py/'
  15. root = './data/cifar10/cifar10/'
  16. header = ['label', 'index']
  17. rows = []
  18. for i in range(5):
  19.     file = path + 'data_batch_' + str(i+1)
  20.     Xtr = unpickle(file)
  21.     print('正在处理训练集：', i)
  22.     for j in range(10000):  # 处理10k张图片
  23.         img = np.reshape(Xtr[b'data'][j], (3, 32, 32))  # R:1024，G:1024，B:1024
  24.         img = img.transpose(1, 2, 0)  # 转置成为图片的格式（H,W,C）
  25.         picName = './data/cifar10/cifar10/train/' + \
  26.             str(int(i*10000+j)) + '.jpg'
  27.         if not os.path.exists(picName):
  28.             imsave(picName, img)
  29.             rows.append((Xtr[b'labels'][j], j))
  30. with open(root+'train/train.csv', 'w', encoding='utf-8', newline='') as f:
  31.     w = csv.writer(f)
  32.     w.writerow(header)
  33.     w.writerows(rows)
  35. rows = []
  36. Xte = unpickle(path+'test_batch')
  37. print('正在处理测试集： 0')
  38. for j in range(10000):  # 处理10k张图片
  39.     img = np.reshape(Xte[b'data'][j], (3, 32, 32))
  40.     img = img.transpose(1, 2, 0)  # 转置成为图片的格式（H,W,C）
  41.     picName = './data/cifar10/cifar10/test/' + \
  42.         str(j) + '.jpg'
  43.     if not os.path.exists(picName):
  44.         imsave(picName, img)
  45.         rows.append((Xte[b'labels'][j], j))
  46. with open(root+'test/test.csv', 'w', encoding='utf-8', newline='') as f:
  47.     w = csv.writer(f)
  48.     w.writerow(header)
  49.     w.writerows(rows)
  50. print('处理完毕')
  1. 正在处理训练集： 0
  2. 正在处理训练集： 1
  3. 正在处理训练集： 2
  4. 正在处理训练集： 3
  5. 正在处理训练集： 4

显示一张图片

  1. from matplotlib.pyplot import imshow
  3. # show a single image
  4. img = np.reshape(Xtr[b'data'][0], (3, 32, 32)).transpose(1, 2, 0)
  5. imshow(img)
  1. <matplotlib.image.AxesImage at 0x18d183d2400>

output_9_1.png

定义自己的类

定义自己的数据集类，主要是在继承（torch.utils.data.Dataset）后修改初始化（__init__(self)）和读取图片（__getitem__(self, index)）的函数（和 __len__(self)）

  • 训练集每张图片的命名规则为“number.jpg”,train.csv标记了各个图片的类别
  • 测试集每张图片的命名规则为“number.jpg”,rain.csv标记了各个图片的类别
  1. import os
  2. import numpy as np
  3. import pandas as pd
  4. import matplotlib.pyplot as plt
  5. from torch.utils.data import Dataset
  6. # import csv
  7. # 数据保存在 ./data/cifar10/cifar10/train(test)/
  8. # root = './data/cifar10/cifar10/
  10. class myDataset(Dataset):
  11.     def __init__(self, root, train=True, transform=None):
  12.         '''
  13.         Args:
  14.             root : 根目录
  15.             transform: 数据处理方式
  16.         '''
  17.         if train:
  18.             self.root = root + 'train/'
  19.             data = pd.read_csv(root + 'train/train.csv')
  20.         else:
  21.             self.root = root + 'test/'
  22.             data = pd.read_csv(root + 'test/test.csv')
  24.         self.imgs = data['index'].values
  25.         self.labels = data['label'].values
  26.         if transform:
  27.             self.transform = transform
  29.     def __getitem__(self, index):
  30.         target = self.labels[index]
  31.         image = plt.imread(self.root+str(int(index))+'.jpg')
  32.         if self.transform is not None:
  33.             image = self.transform(image)
  34.         return image, target
  36.     def __len__(self):
  37.         return len(self.imgs)

使用自己的类

  1. from torch.utils.data import Dataset, DataLoader
  3. batch_size = 4
  4. epoch_size = 3
  6. transform = transforms.Compose(
  7.     [transforms.ToTensor(),
  8.      transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])
  10. trainset = myDataset(root='./data/cifar10/cifar10/', 
  11.                      train=True,
  12.                      transform=transform)
  13. trainloader = torch.utils.data.DataLoader(trainset,
  14.                                         batch_size=batch_size,
  15.                                         shuffle=True,
  16.                                         num_workers=0)
  19. testset = myDataset(root='./data/cifar10/cifar10/', 
  20.                     train=False,
  21.                     transform=transform)
  22. testloader = torch.utils.data.DataLoader(testset,
  23.                                         batch_size=batch_size,
  24.                                         shuffle=False,
  25.                                         num_workers=0)
  27. classes = ('plane', 'car', 'bird', 'cat',
  28.            'deer', 'dog', 'frog', 'horse', 'ship', 'truck')

BUG

当 shuffle 设置为 True 时，num_sample 报错

原因

因为之前的csv文件被清空导致数据为空

  1. def imshow(img):
  2.     img = img / 2 + 0.5  # 去标准化
  3.     npimg = img.numpy()  # 将torch.FloatTensor 转换为numpy
  4. #     plt.axis("off")  # 不显示坐标尺寸
  5.     plt.imshow(np.transpose(npimg, (1, 2, 0)))  # 进行转置
  6.     plt.show()  # 显示图片
  9. # get some random training images
  10. dataiter = iter(trainloader)
  11. images, labels = dataiter.next()
  13. # show images
  14. imshow(torchvision.utils.make_grid(images))
  15. # print labels
  16. print(' '.join('%11s' % classes[labels[j]] for j in range(batch_size)))

output_16_0.png

  1. plane         dog       truck        frog

torchvision.datasets.ImageFolder

保存/加载模型

保存模型

  1. # 模型保存
  2. torch.save(net.state_dict(), './model/learn0.pt')  # .pt or .pth

定义网络

  1. # 网络结构
  2. class Net(nn.Module):  # nn.Module 是所有神经网络的基类，自定义的网络应该继承自它
  3.     def __init__(self):
  4.         super(Net, self).__init__()
  5.         self.conv1 = nn.Conv2d(3, 6, 5)
  6.         self.pool = nn.MaxPool2d(2, 2)
  7.         self.conv2 = nn.Conv2d(6, 16, 5)
  8.         self.fc1 = nn.Linear(16 * 5 * 5, 120)
  9.         self.fc2 = nn.Linear(120, 84)
  10.         self.fc3 = nn.Linear(84, 10)
  12.     def forward(self, x):
  13.         x = self.pool(F.relu(self.conv1(x)))
  14.         x = self.pool(F.relu(self.conv2(x)))
  15.         x = x.view(-1, 16 * 5 * 5)
  16.         x = F.relu(self.fc1(x))
  17.         x = F.relu(self.fc2(x))
  18.         x = self.fc3(x)
  19.         return x

加载模型

  1. model = Net().to('cuda')
  2. model.load_state_dict(torch.load('./model/learn0.pt'))
  3. model.eval()
  1. Net(
  2.   (conv1): Conv2d(3, 6, kernel_size=(5, 5), stride=(1, 1))
  3.   (pool): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
  4.   (conv2): Conv2d(6, 16, kernel_size=(5, 5), stride=(1, 1))
  5.   (fc1): Linear(in_features=400, out_features=120, bias=True)
  6.   (fc2): Linear(in_features=120, out_features=84, bias=True)
  7.   (fc3): Linear(in_features=84, out_features=10, bias=True)
  8. )

model.load_state_dict(torch.load(‘./model/learn0.pt’))

  1. IncompatibleKeys(missing_keys=[], unexpected_keys=[])

猜测是加载模型正确，没有遗漏和错误。

测试模型

  1. # 加载测试集
  2. test_set = torchvision.datasets.CIFAR10(
  3.     root='./data/cifar10',
  4.     train=False,
  5.     transform=transforms.Compose(
  6.         [transforms.ToTensor(),
  7.          transforms.Normalize((0.5, 0.5, 0.5),
  8.                               (0.5, 0.5, 0.5))])
  9. )
  10. testloader = torch.utils.data.DataLoader(
  11.     test_set,
  12.     batch_size=4,
  13.     shuffle=True,
  14.     num_workers=2
  15. )
  16. classes = ('plane', 'car', 'bird', 'cat',
  17.            'deer', 'dog', 'frog', 'horse', 'ship', 'truck')
  1. # 测试
  2. dataiter = iter(testloader)
  3. images, labels = dataiter.next()
  6. def imshow(img):
  7.     img = img / 2 + 0.5  # 去标准化
  8.     npimg = img.numpy()  # 将torch.FloatTensor 转换为numpy
  9. #     plt.axis("off")  # 不显示坐标尺寸
  10.     plt.imshow(np.transpose(npimg, (1, 2, 0)))  # 进行转置
  11.     plt.show()  # 显示图片
  14. # print images
  15. print('GroundTruth:')
  16. imshow(torchvision.utils.make_grid(images))
  17. print(' '.join('%11s' % classes[labels[j]] for j in range(4)))
  19. # 测试一次
  20. outputs = model(images.to('cuda'))
  21. _, predicted = torch.max(outputs, 1)
  22. print('Predicted:\n', ' '.join('%11s' % classes[predicted[j]]
  23.                                for j in range(4)))
  1. GroundTruth:

output_28_1.png

  1. dog        deer         dog        bird
  2. Predicted:
  3.          cat         car         cat        bird

测试自己的图片

网络输入尺寸是 32x32，因此需要把自己的图片进行预处理，之后才可以进行测试

需要注意的是，如果使用 plt.imread() 读入图片，返回的是 np.ndarray 类型，此时需要用 transforms.ToPILImage() 转化为 PIL 类型

采用 PIL.Image(Pillow for Python3.x) 则无需进行转化

  1. img_path = './data/cifar10/mytest/car-1.jpg'  # √
  2. # img_path = './data/cifar10/mytest/dog-1.jpg'  # ×
  4. # # 采用 plt.imread() 读入图片
  5. # img = plt.imread(img_path)
  6. # print('图片尺寸:',np.array(img).shape,type(img))
  7. # transform = transforms.Compose(
  8. #     [transforms.ToPILImage(), 
  9. #      transforms.Resize([32,32]),
  10. #      transforms.ToTensor(),
  11. #      transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])
  13. # 采用 PIL.Image 读取图片
  14. from PIL import Image
  15. img = Image.open(img_path)
  16. print('未 transform 的图片尺寸:',np.array(img).shape,'  图片类型：',type(img))
  17. plt.imshow(np.array(img))
  18. plt.show()
  20. transform = transforms.Compose(
  21.      [transforms.Resize([32,32]),
  22.      transforms.ToTensor(),
  23.      transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])
  26. img = transform(img).reshape(-1,3,32,32)
  27. print('已 transpose 的图片尺寸:',np.array(img).shape,'  图片类型：',type(img))
  28. # plt.imshow(img)
  29. imshow(torchvision.utils.make_grid(img))
  30. output = model(img.to('cuda'))
  31. _, predicted = torch.max(output, 1)
  32. print('模型输出为：',classes[predicted])
  1. 未 transform 的图片尺寸: (435, 728, 3)   图片类型： <class 'PIL.JpegImagePlugin.JpegImageFile'>

output_31_1.png

  1. 已 transpose 的图片尺寸: (1, 3, 32, 32)   图片类型： <class 'torch.Tensor'>

output_31_3.png

  1. 模型输出为： car
  1. !pip install Pillow
  1. Looking in indexes: https://pypi.doubanio.com/simple/
  2. Requirement already satisfied: Pillow in d:\office\office\python\python3.6.2\lib\site-packages (4.3.0)
  3. Requirement already satisfied: olefile in d:\office\office\python\python3.6.2\lib\site-packages (from Pillow) (0.44)

格式转换

  1. # !jupyter nbconvert --to html --template full learn.ipynb
  2. # !jupyter nbconvert --to markdown learn.ipynb
  3. !jupyter nbconvert --to html --template full learn-1.ipynb
  4. !jupyter nbconvert --to markdown learn-1.ipynb
  1. [NbConvertApp] Converting notebook learn-1.ipynb to html
  2. [NbConvertApp] Writing 335545 bytes to learn-1.html
  3. [NbConvertApp] Converting notebook learn-1.ipynb to markdown
  4. [NbConvertApp] Support files will be in learn-1_files\
  5. [NbConvertApp] Making directory learn-1_files
  6. [NbConvertApp] Making directory learn-1_files
  7. [NbConvertApp] Writing 8365 bytes to learn-1.md