Pytorch02:使用自己的数据搭建VGG16网络
写在前面
在Pytorch01:使用标准数据集CIFAR-10搭建VGG16网络
中介绍了如何利用标准的数据集类来搭建一个简单的VGG-16网络。然而在实际应用中，根据需要使用自己的数据集才是常见的情况。因此，本文将介绍如何在Pytorch中用自己的数据构建数据集类以替换标准的CIFAR-10等数据集，作为输入供神经网络运行。
一、要实现什么东西
首先我们要了解，如果是使用自己的数据，代码部分和之前有什么不同。如果是使用标准的数据集类，如CIFAR-10，可以用以下的代码完成数据的输入和加载：
# 下载训练集 CIFAR-10训练集
train_dataset = datasets.CIFAR10('./data', train=True, transform=transforms.ToTensor(), download=True)
train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
test_dataset = datasets.CIFAR10('./data', train=False, transform=transforms.ToTensor(), download=True)
test_loader = DataLoader(test_dataset, batch_size=32, shuffle=False)
其中，训练数据和测试数据均分为两步进行处理：第一步，通过Dataset类将外部的数据文件读取进来，形成数据的接口；第二步，使用DataLoader来进一步封装数据，以满足神经网络的批量输入等需求。而个人自定义的部分主要是实现第一步中的类似于datasets.CIFAR10的数据集类，并在其中完成数据读取和处理的相关过程。
二、所使用的数据集
本文处理的数据集是一个从社交媒体上爬取的图片集，用于描述各种灾难的场景。数据集下载地址为：Multimodal Damage Identification for Humanitarian Computing Data Set。数据集的形式如下图，在图片的名字中含有其标签，所有图片放在同一个文件夹下。为了简便，本文只处理一个文件夹中的图片。
根据数据集中含有的图片标签，我们首先定义标签的编号，供后续读取图片标签时使用。
label_set = {
'accrafloods': 0,
'buildingcollapse': 1,
'destroyedbuilding': 2,
'destruction': 3,
'disasters': 4,
'earthquake': 5,
'earthquakenepal': 6,
'floodwater': 7,
'hurricaneharvey': 8,
'hurricaneirma': 9,
'hurricanematthew': 10,
'hurricanesandy': 11,
'isiscrimes': 12,
'naturaldisaster': 13,
'sandydamage': 14,
'suicidebombing': 15,
'syriawarcrimes': 16,
'terrorattack': 17,
'warsyria': 18,
'wreckedcar': 19,
'yemencrisis': 20,
'buildingfire': 21,
}
三、自定义数据集类
自定义数据集类需要继承torch.utils.data.Dataset类，然后重写它的__len__和__getitem__函数。一个示例的自定义数据集类定义如下：
class CrisisMMDDataset(Dataset):
def __init__(self, root_dir, transform=None):
"""
:param root_dir: 图片文件夹路径
:param transform: 图片转Tensor的操作
"""
self.root_dir = root_dir
self.transform = transform
self.picname_list = os.listdir(root_dir) # 列出root_dir下的所有文件
# 返回数据集长度，为重载函数
def __len__(self):
return len(self.picname_list)
# 获得一个样本，为重载函数
def __getitem__(self, idx):
# print(self.picname_list[idx])
img_name = os.path.join(self.root_dir, self.picname_list[idx]) # 拼接图片的路径
image = io.imread(img_name) # skimage.io.imread返回的直接是numpy.ndarray类型
# 文件名示例：accrafloods_2015-06-04_22-54-13
label = label_set[self.picname_list[idx].split('_')[0]] # 返回对应标签名称对应的编号
label = np.array(label) # 转为numpy.ndarray类型
sample = {
'image': image, 'label': label} # 使用字典结构封装图片和对应的标签
if self.transform:
sample = self.transform(sample) # transform是转换图片数据函数
return sample
__len__的作用是返回数据集的总大小，即含有多少个<样本-标签>对。__getitem__的作用是获得一个<样本-标签>对，其中，sample的格式可以自己定义，可以返回一个元组tuple类型，也可以返回一个字典类型或者其它的类型，本文采用的是字典类型（主要是增加了对样本和标签的命名以显得更为清晰）。transform=的传入参数可以是torchvision.transforms下的标准转换对象，即transforms.ToTensor()，用于将图片转换成Tensor类型；或者是transforms.ToTensor()和其它的转换对象，如transforms.Scale()，用于缩放图片的尺寸等，的组合。但是要注意的是，标准的transforms对象只能转换图片为Tensor类型，且要求输入的图片为PIL Image或者numpy.ndarray类型。因此在上述代码里，图片和标签均转换成numpy.ndarray类型，这一方面是为了能够满足标准transforms对象的输入类型要求，另一方面也是因为从numpy.ndarray类型能够较为方便地转为Tensor类型。另外，如果是要使用标准的transforms对象的话，还需要单独将标签也转为Tensor类型。- 除了使用标准的
transforms对象，也可以使用自己定义的转换对象。本文自定义了如下的转换类，用于统一将图片和标签一起转换成Tensor类型。总之，最后返回的sample一定是Tensor类型的。
# 缩放图片
class Rescale(object):
def __init__(self, output_size=32):
assert isinstance(output_size, (int, tuple)) # 限制输入的尺寸为int或者tuple类型
self.output_size = output_size
def __call__(self, sample):
image, label = sample['image'], sample['label']
h, w = image.shape[:2] # [:2]只取0和1，2不取
if isinstance(self.output_size, int):
new_h, new_w = self.output_size, self.output_size
else:
new_h, new_w = self.output_size
new_h, new_w = int(new_h), int(new_w)
img = transform.resize(image, (new_h, new_w))
return {
'image': img, 'label': label}
# 将样本数据转为Tensors类型
class ToTensor(object):
def __call__(self, sample):
image, label = sample['image'], sample['label']
# 交换颜色轴因为
# numpy包的图片是: H * W * C
# torch包的图片是: C * H * W
image = image.transpose((2, 0, 1))
# 将numpy数组转换成tensor类型，才能给神经网络使用
image = torch.from_numpy(image)
label = torch.from_numpy(label)
# 将图片的tensor转为FloatTensor
image = image.type(torch.FloatTensor)
label = label.type(torch.LongTensor)
return {
'image': image,
'label': label}
四、使用自定义数据集类生成对象
类似于标准的CIFAR-10数据集对象，我们现在可以调用上面定义的类来生成自定义的数据集对象train_dataset和test_dataset了。其中，transforms.Compose的作用就是依次调用其传入的对象处理数据。
pic_size = 32 # 图片尺寸
train_dataset = CrisisMMDDataset(root_dir='./data/CrisisMMD/resize_data',
transform=transforms.Compose([Rescale(pic_size), ToTensor()]))
test_dataset = CrisisMMDDataset(root_dir='./data/CrisisMMD/resize_flip_data',
transform=transforms.Compose([Rescale(pic_size), ToTensor()]))
关于训练集和测试集：
- 可以像上述代码一样，传入不同的数据文件夹分别产生训练集和测试集。如果是需要只传入一个文件夹，并将其中的数据部分作为训练集，部分作为测试集的话，可以用下面的代码来划分，主要是使用了
torch.utils.data.random_split函数：
# 使用统一一个数据集后分割成测试集和测试集
whole_dataset = CrisisMMDDataset(root_dir='./data/CrisisMMD/resize_data',
transform=transforms.Compose([Rescale(pic_size), ToTensor()]))
# 分割训练集和测试集
train_size = int(0.8 * len(whole_dataset)) # 训练集占0.8
test_size = len(whole_dataset) - train_size # 测试集占0.2
# torch.utils.data.random_split函数用于随机分割数据集
train_dataset, test_dataset = random_split(whole_dataset, [train_size, test_size])
- 甚至可以通过在
CrisisMMDDataset中根据传入参数返回不同的数据，如标准的CIFAR-10数据集的训练集和测试集是通过设置train=的值来控制数据集的生成，但这要求更加复杂的数据集类定义，本文就不进行进一步的实现了。
五、通过DataLoader生成批量数据
从数据集类对象中获取到的数据都是一条一条的<样本-标签>对，然而，神经网络的输入通常要求是批量输入的，这就需要使用DataLoader来生成一个batch一个batch的数据。代码如下：
train_loader = DataLoader(train_dataset,
batch_size=batch_size,
shuffle=True,
num_workers=0,
pin_memory=True)
test_loader = DataLoader(test_dataset,
batch_size=batch_size,
num_workers=0,
pin_memory=True)
shuffle=True是在每个epoch开始的时候，对数据进行重新排序以打乱数据的顺序。- 使用
num_workers要放在main函数中，不然会报错。num_workers决定了有几个进程来处理data loading的过程，0意味着所有的数据都在主进程处理，默认为0。 pin_memory=True即将读取到的数据都放在内存中而非虚拟内存，如果你的电脑内存够大，可以用它来提高I/O速度。
六、训练过程和测试过程
for epoch in range(num_epoches):
model.train() # 训练开始
print('*' * 25, 'epoch {}'.format(epoch + 1), '*' * 25)
running_loss = 0.0
running_acc = 0.0
for i, data in enumerate(train_loader, 1): # 注意不要用错了train_dataset
img, label = data['image'], data['label'] # 自定义数据集
# cuda
if use_gpu:
img = img.cuda()
label = label.cuda()
# 向前传播
out = model(img)
loss = criterion(out, label)
running_loss += loss.item() * label.size(0)
_, pred = torch.max(out, 1)
num_correct = (pred == label).sum()
accuracy = (pred == label).float().mean()
running_acc += num_correct.item()
# 向后传播
optimizer.zero_grad()
loss.backward()
optimizer.step()
print('Finish {} epoch, Loss: {:.6f}, Acc: {:.6f}'.format(
epoch + 1, running_loss / (len(train_dataset)), running_acc / (len(train_dataset))))
model.eval() # 训练结束，测试开始
eval_loss = 0
eval_acc = 0
with torch.no_grad():
for data in test_loader: # 注意不要用错了test_dataset
img, label = data['image'], data['label'] # 自定义数据集
if use_gpu:
img = img.cuda()
label = label.cuda()
out = model(img)
loss = criterion(out, label)
eval_loss += loss.item() * label.size(0)
_, pred = torch.max(out, 1)
num_correct = (pred == label).sum()
eval_acc += num_correct.item()
print('Test Loss: {:.6f}, Acc: {:.6f}'.format(
eval_loss / len(test_dataset), eval_acc / len(test_dataset)))
和之前的唯一不同在于，改用了字典读取的方式来获取自定义数据集中的一个Batch的样本数据。
七、关于读取数据时的I/O瓶颈
该数据集中的图片尺寸均在640*640左右，在我的电脑上读取数据时就出现了I/O瓶颈，具体表现为GPU的CUDA在大部分时间内不工作，GPU利用率极低。这是因为GPU需要等CPU读取数据形成一个Batch之后再传入GPU运行，因此出现了I/O瓶颈。相比之下，标准的CIFAR-10数据集在处理时CUDA几乎满负荷运行，GPU的利用率极高。
简单的解决办法有三个：
- 使用
DataLoader时将num_workers的数量增大。这个主要是通过利用CPU的多核多进程处理数据，降低I/O瓶颈。一般而言，增加num_workers能够让GPU的高利用率保持长一点的时间，但无法从根本上避免GPU在大部分时间中不工作的现象。 - 对图片进行预处理，例如减小图片的尺寸。由于标准的CIFAR-10数据集原始图片为
32*32，它的每一张图片的读取速度快很多。因此，我们也可以提前将图片的尺寸缩小，如在本文中，所使用的VGG-16神经网络的输入仍是32*32，所以可以将缩小图片尺寸的工作提前进行。使用以下代码可以将原数据集转换成一个32*32大小的图片数据集。这相当于将原来自定义数据集类中的Rescale()过程提前做了。通过预处理，能够极大地提高GPU的利用率，达到和使用标准CIFAR-10数据集类似的效果。
# 预处理：更改数据集图片大小并保存在另一个文件夹下
root_dir = './data/CrisisMMD/damaged_infrastructure/images' # 原数据集路径
save_dir = './data/CrisisMMD/resize_data' # resize后数据集保存路径
output_size = 32
picname_list = os.listdir(root_dir)
for picname in picname_list:
img_name = os.path.join(root_dir, picname)
print(img_name)
image = io.imread(img_name)
# 对输入的图片尺寸类型，含int或者tuple，做统一化处理
h, w = image.shape[:2] # [:2]只取0和1，2不取
if isinstance(output_size, int):
new_h, new_w = output_size, output_size
else:
new_h, new_w = output_size
new_h, new_w = int(new_h), int(new_w)
# 更改图片尺寸
img = transform.resize(image, (new_h, new_w))
# 保存图片
save_path = os.path.join(save_dir, picname)
io.imsave(save_path, img)
- 另外，还可以考虑提前将图片转换成类似于标准CIFAR-10数据集的二进制文件，也可以提高读取的速度，降低I/O瓶颈。可以参考：用自己的数据，制作python版本的cifar10数据集，这里就不再进一步说明了。
图片的resize操作如果能够在网络之前预处理，尤其是当图片过大需要转换成小图片的时候，能够极大地提高网络的io性能瓶颈，增加GPU利用率。大图片对于CPU的处理性能要求很高，不然会卡在__getitem__函数里面，CUDA使用率大部分时间为0，偶尔尖峰跳跃。小图片时，num_workers=0即只用一个CPU，没有了并行的处理过程，反而能够提高GPU的CUDA使用率，但是此时CPU的使用率很高。然而如果增加num_workers数量，CPU使用率能够下降，但GPU的CUDA使用率会尖峰跳跃，无法拉满。
在这篇博客中有关于GPU利用率低的其他情况较为详尽的说明，可供进一步参考。
八、修改网络以适应不同的图片输入尺寸
在上述代码中，我们仍使用32*32的输入图片尺寸，但是由于原始的图片是640*640的，在过度压缩图片的过程中必然损失了大量的信息，导致最后训练的准确率不高。那我们能不能修改网络的输入图片尺寸，比如改成64*64呢？当然可以，主要修改两个地方就可以了。
- 修改网络的第一个全连接层的输入尺寸。这是因为卷积层和池化层的定义和使用并不受输入图片的影响，但是全连接层的定义需要提前设定其输入的尺寸，这是受图片大小影响的。如果原始图片输入尺寸为
32*32，那么在最后一个池化层输出是512*1*1的张量，所以全连接层的输入是512。若改变网络的输入尺寸，可以增加input_size变量并修改网络为：
class VGG16(nn.Module):
def __init__(self, input_size, num_classes=22):
super(VGG16, self).__init__()
self.features = nn.Sequential(
# 1, 224*224*3, conv=3*3*64
nn.Conv2d(3, 64, kernel_size=3, padding=1),
nn.BatchNorm2d(64),
nn.ReLU(True),
...
# 13
nn.Conv2d(512, 512, kernel_size=3, padding=1),
nn.BatchNorm2d(512),
nn.ReLU(True),
nn.MaxPool2d(kernel_size=2, stride=2),
)
self.classifier = nn.Sequential(
# 14, 512=>4096
nn.Linear(512*int(input_size/32)*int(input_size/32), 4096), # 修改输入的尺寸
nn.ReLU(True),
nn.Dropout(),
# 15, 4096=>4096
nn.Linear(4096, 4096),
nn.ReLU(True),
nn.Dropout(),
# 16, 4096=>output_sizes
nn.Linear(4096, num_classes),
)
def forward(self, x):
...
# 创建model实例对象
model = VGG16(pic_size)
- 修改数据读取时的转换函数。即对应修改
Rescale(pic_size)中的参数值，这将决定是不是真的修改了输入到网络中的图片尺寸大小。注意，原始图片的大小并不是输入网络中的图片大小。决定图片大小的是读取图片过程中的处理函数。
这里默认修改的图片尺寸为32的倍数，如64, 128等。如果不是32的倍数，则不能按照上面的int(input_size/32)来修改，而是要通过计算（即每次2*2池化图片的长和宽均减半）来得到最后的池化输出的图片大小。要是实在不会算，或者怕算得不准确，也可以通过在forward中直接打印输出尺寸，如下：
def forward(self, x):
# x.size=batch_size, channels, width, height
# 如果输入尺寸为[64, 3, 64, 64]，即图片尺寸为64*64
out = self.features(x)
# 输出最后一次池化后的张量尺寸
print(out.size(0), out.size(1), out.size(2), out.size(3))
# 那么打印的结果为[64, 512, 2, 2]，即第一个全连接层的输入是512*2*2
out = out.view(out.size(0), -1)
out = self.classifier(out)
return out
九、一些错误
- 数据和标签转换Tensor类型的错误。
神经网络中的参数类型是torch.FloatTensor类型，在转换的时候，图片和标签直接调用torch.from_numpy转换的是torch.DoubleTensor，输入到神经网络中就会出现与参数类型不匹配的错误。另外，由于标签是不输入到神经网络中的，所以并不会在神经网络中出现和参数类型不匹配的错误，但会在调用criterion的时候出现和输出向量类型不匹配的错误。
解决方法：增加一行用于类型转换，将torch.DoubleTensor转换成torch.FloatTensor和torch.LongTensor即可。
image = image.type(torch.FloatTensor)
label = label.type(torch.LongTensor)
- OMP: Error #15: Initializing libiomp5md.dll, but found libiomp5md.dll already initialized
这个问题是由于导入的包之间存在相同的.dll文件导致的。
解决方法：最为保险的方式就是在代码开头增加两行代码。可参考这篇博客。
import os
os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"