第一步：准备图片

首先我们需要准备好用于训练和验证的图片，图片的分辨率可以随意，因为pytorch有resize的功能，可以把图片缩放到任意你想要的大小。图片支持类型：.JPG，.PNG…（反正常用的一般都支持）
详细参考官网：
https://pillow.readthedocs.io/en/5.1.x/handbook/image-file-formats.html

第二步：图片分类

我们需要创建一个文件夹（最好是和代码文件在同级目录下，虽然我没有这样做哈哈），你可以取名为test啊什么的，反正随意啦。接着我们要在这个文件夹下再创建两个文件夹，分别命名为train和valid，存放训练数据和验证数据，哈哈别急还没完。我们还需要在这两个文件夹下创建n个文件夹（禁止套娃！），文件夹的个数取决你的分类的类别，有多少类就要创建多少个文件夹，一个文件夹对应一个类。你也可以看下面的导图。创建完后就可以把你准备好的图片按照对应的文件夹进行存放啦

文件夹分类导图

第三步：图片导入

完整代码：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
#常用头文件
from glob import glob
import os
import numpy as np
from PIL import Image
import matplotlib.pyplot as plt
import shutil
from torchvision import transforms
from torchvision import models
import torch
from torch.autograd import Variable
import torch.nn as nn
from torch.optim import lr_scheduler
from torch import optim
from torchvision.datasets import ImageFolder
from torchvision.utils import make_grid
import time
import cv2

train = ImageFolder('C:/Users/Arthur/Desktop/sum/tempu/train/',simple_transform)
valid = ImageFolder('C:/Users/Arthur/Desktop/sum/tempu/valid/',simple_transform)
train_data_gen = torch.utils.data.DataLoader(train,shuffle=True,batch_size=16,num_workers=0,drop_last=True)
valid_data_gen = torch.utils.data.DataLoader(valid,shuffle=True,batch_size=16,num_workers=0,drop_last=True)
dataset_sizes = {'train':len(train_data_gen.dataset),'valid':len(valid_data_gen.dataset)}
dataloaders = {'train':train_data_gen,'valid':valid_data_gen}

注：

  1. 如果你的代码文件和你的数据集在同一文件夹下，那么引用的时候就可以不用全路径，如：
1
2
train = ImageFolder('/tempu/train/',simple_transform)
valid = ImageFolder('/tempu/valid/',simple_transform)
  1. simple_transform是你需要应用的数据变换和增强
  2. shuffle=True：训练时随机抽取样本
  3. drop_last=True：当你的数据集不能被定义的batch_size整除时，多余的部分会被丢掉。
  4. num_workers=0 启用多线程个数，一般设置成0是不会有任何问题的。启用多线程可以加快数据集读取速度，但是很多时候需要自己调试，有些IDE并不支持。