不管是区分外星人和铁血战士，还是对图像进行分类的其他问题，都可以利用机器学习的方法来进行处理，而且使用起来更高效、更有趣。
话不多说，本文利用PyTorch这一深度学习框架，来对外星人和铁血战士进行识别分类。主要包括以下几个部分：
- 准备数据集
- 导入相关依赖
- 创建数据生成器
- 创建网络
- 训练模型
- 保存并加载模型
- 对样本测试图像进行预测
1.准备数据集
数据集无疑就是外星人和铁血战士的图像了，这里使用的都是大约 255×255 像素的图像，你可以从 Kaggle 直接下载该数据集。
该数据集包含两部分：
- 训练集——包括347张外星人图像，347张铁血战士图像，该训练集用于网络模型的训练。
- 验证集——包括100张外星人图像，100张铁血战士图像，该验证集用于检测训练好的模型在以前未见过的数据上的性能好坏(也就是说，用之前训练好的模型来识别一下新的图像，看看该模型训练的效果怎么样)
数据集结构如下所示：
2.导入相关依赖
数据集准备好了之后，先把它放到一边，我们来准备一下用到的环境。
文中用到了Keras和PyTorch这两种深度学习框架，都可以实现对图片的分类，我这里就使用PyTorch来进行本次的实验吧。
本次实验环境如下所示：
- Python-3.6，只要是3.5版本以上都可以
- PyTorch-1.0.1
- torchvision-0.2.2
- jupyter-1.0.1
安装完之后就可以导入相应的模块，在终端运行jupyter notebook demo.ipynb后就可以键入代码了。如下所示：
1
2
3
4
5
6
7
8
9
|
import numpy as np
import matplotlib.pyplot as plt
from PIL import Image
import torch
from torchvision import datasets, models, transforms
import torch.nn as nn
from torch.nn import functional as F
import torch.optim as optim
|
当然，你也可以查看一下模块的版本：
1
|
torch.__version__ # 输出'1.0.1'
|
1
2
|
import torchvision
torchvision.__version__ # 输出'0.2.2'
|
3.创建数据生成器
通常情况下，图像不能一次全部加载完，因为这样会占用太多内存空间。因此，这里使用数据生成器批量加载图像，比如说，每次加载32张图像。每次遍历整个数据集的过程称为一个 epoch。
于此同时，使用数据生成器对图像进行预处理，调整图像大小并标准化图像。因为这里使用的是 ResNet-50 网络，所以要将图片处理成 244×244 像素，带有缩放的颜色通道的图像。此外，在这次实验中，还对图像进行了随机剪切，缩放和水翻转。
以下是创建数据生成器的过程：
- 从文件中加载数据
- 规范化数据（包括训练和验证）
- 增加数据（仅用于训练）
代码如下：
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
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
|
input_path = "./data/"
normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
data_transforms = {
'train':
transforms.Compose([
transforms.Resize((224,224)),
transforms.RandomAffine(0, shear=10, scale=(0.8,1.2)),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
normalize
]),
'validation':
transforms.Compose([
transforms.Resize((224,224)),
transforms.ToTensor(),
normalize
]),
}
image_datasets = {
'train':
datasets.ImageFolder(input_path + 'train', data_transforms['train']),
'validation':
datasets.ImageFolder(input_path + 'validation', data_transforms['validation'])
}
dataloaders = {
'train':
torch.utils.data.DataLoader(image_datasets['train'],
batch_size=32,
shuffle=True,
num_workers=0), # for Kaggle
'validation':
torch.utils.data.DataLoader(image_datasets['validation'],
batch_size=32,
shuffle=False,
num_workers=0) # for Kaggle
}
|
4.创建网络
接下来是导入预先训练好的ResNet-50模型，这里仅训练最后两个完全连接（密集）层，并调整最后一层。
因为我的笔记本显卡不支持PyTorch的GPU版本，所以就用CPU来跑的代码。
步骤如下：
- 加载预先训练好的网络，切断头部并冻结权重
- 添加自定义的密集层（这里的隐层选择了128个神经元）
- 设置优化器和损失函数
代码如下：
1
2
3
4
5
6
7
8
9
10
11
12
13
|
device = torch.device("cuda:0" if not torch.cuda.is_available() else "cpu")
model = models.resnet50(pretrained=True).to(device)
for param in model.parameters():
param.requires_grad = False
model.fc = nn.Sequential(
nn.Linear(2048, 128),
nn.ReLU(inplace=True),
nn.Linear(128, 2)).to(device)
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.fc.parameters())
|
5.训练模型
现在可以进行模型的训练了，我们需要传递数据，计算其损失函数并修改相应的网络权重。步骤如下：
- 训练模型
- 测量训练集和验证集的损失函数（对数损失）和准确性
代码如下：
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
26
27
28
29
30
31
32
33
34
35
36
37
38
39
|
def train_model(model, criterion, optimizer, num_epochs=3):
for epoch in range(num_epochs):
print('Epoch {}/{}'.format(epoch+1, num_epochs))
print('-' * 10)
for phase in ['train', 'validation']:
if phase == 'train':
model.train()
else:
model.eval()
running_loss = 0.0
running_corrects = 0
for inputs, labels in dataloaders[phase]:
inputs = inputs.to(device)
labels = labels.to(device)
outputs = model(inputs)
loss = criterion(outputs, labels)
if phase == 'train':
optimizer.zero_grad()
loss.backward()
optimizer.step()
_, preds = torch.max(outputs, 1)
running_loss += loss.item() * inputs.size(0)
running_corrects += torch.sum(preds == labels.data)
epoch_loss = running_loss / len(image_datasets[phase])
epoch_acc = running_corrects.double() / len(image_datasets[phase])
print('{} loss: {:.4f}, acc: {:.4f}'.format(phase,
epoch_loss,
epoch_acc))
return model
model_trained = train_model(model, criterion, optimizer, num_epochs=3)
|
输出结果如下所示：
1
2
3
4
5
6
7
8
9
10
11
12
|
Epoch 1/3
----------
train loss: 0.6031, acc: 0.6888
validation loss: 0.4098, acc: 0.8700
Epoch 2/3
----------
train loss: 0.3978, acc: 0.8357
validation loss: 0.3342, acc: 0.8700
Epoch 3/3
----------
train loss: 0.2746, acc: 0.9020
validation loss: 0.2431, acc: 0.9200
|
6.保存并加载模型
保存
训练一个模型通常需要很长的时间和计算成本，最好的方法是将其保存下来以便日后使用。通常有如下两种方法：
- 将整个模型的体系结构和训练权重（以及优化器状态）保存到文件中
- 将训练过的权重保存到文件中（将模型体系结构保留在代码中）
代码如下：
1
|
torch.save(model_trained.state_dict(), './models/pytorch/weights.h5')
|
加载（导入模型）
1
2
3
4
5
6
|
model = models.resnet50(pretrained=False).to(device)
model.fc = nn.Sequential(
nn.Linear(2048, 128),
nn.ReLU(inplace=True),
nn.Linear(128, 2)).to(device)
model.load_state_dict(torch.load('./models/pytorch/weights.h5'))
|
对于保存和加载模型这部分，还需要注意一下文件的路径是否准确。
7.对样本测试图像进行预测
步骤如下：
- 加载和预处理测试图像
- 预测图像类别
- 显示图像以及预测
示例代码：
1
2
3
4
5
6
7
|
validation_img_paths = ["validation/alien/12.jpg",
"validation/alien/23.jpg",
"validation/predator/34.jpg"]
img_list = [Image.open(input_path + img_path) for img_path in validation_img_paths]
validation_batch = torch.stack([data_transforms['validation'](img).to(device)
for img in img_list])
|
这里可以打印一些信息便于查看：
1
2
|
pred_logits_tensor = model(validation_batch)
pred_logits_tensor
|
输出结果如下：
1
2
3
|
tensor([[ 0.2340, -0.3168],
[ 0.9597, -0.9461],
[-2.0024, 1.4126]], grad_fn=<AddmmBackward>)
|
同时可以查看以下信息：
1
2
|
pred_probs = F.softmax(pred_logits_tensor, dim=1).cpu().data.numpy()
pred_probs
|
输出结果如下：
1
2
3
|
array([[0.63430333, 0.36569664],
[0.8705474 , 0.12945262],
[0.03182863, 0.9681714 ]], dtype=float32)
|
最后的分类显示结果：
1
2
3
4
5
6
|
fig, axs = plt.subplots(1, len(img_list), figsize=(20, 5))
for i, img in enumerate(img_list):
ax = axs[i]
ax.axis('off')
ax.set_title("{:.0f}% Alien, {:.0f}% Predator".format(100*pred_probs[i,0], 100*pred_probs[i,1]))
ax.imshow(img)
|
最终输出的结果：
总结
这原本是由 这篇文章 引发的一个深度学习框架选择的问题，用 Keras 还是 PyTorch。Keras 用起来更加简洁，但代价是它的灵活性；而 PyTorch 提供了更加明确和详细的代码，在大多数情况下，它有着可调试以及灵活的代码，然而在训练时，PyTorch 稍显得有些笨重。至于选择什么样的学习框架，你可以自行阅读该文章，选择一个适合自己的框架进行学习。
参考
Keras vs PyTorch: how to distinguish Aliens vs Predators with transfer learning