PyTorch 简单分类模型示例

之前实习的时候 PyTorch 用的非常多。PyTorch 的优势是非常灵活,语法非常 pythonic,很接近 NumPy。动态图的特性也让调试变得比 TensorFlow 轻松许多。
不过对于作业 / Kaggle 中相对简单的分类任务,Keras 还是我的第一选择:开箱即用,并且和 scikit-learn 有着类似的接口,因此只需少量的代码即可实现模型的训练和迭代。而在 PyTorch 中,很多相似的功能都需要自己手动多写几行代码,对我这个懒人不太友好。
在帮朋友解决问题的时候,发现对 PyTorch 有些生疏,而且官方的例子里并没有具体的例子讲述如何用 NumPy 的数据作为数据集,于是这里便记录一下,方便以后参考。

导入库

第一步是需要导入各种库

1
2
3
4
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim

数据集加载

比较 tricky 的就是构建 PyTorch 中的 Dataset 和 DataLoader 了。虽然非常好用,比如可以自定义实现 DataLoader、Transform 等,但刚上手的时候容易出错。
在用 train_test_split 将 NumPy 数据集划分好后,第一步需要将它们转换为 PyTorch 中的 Tensor,之后用 TensorDataset 创建 Dataset 对象,最后用 DataLoader 创建数据集加载器,之后在训练中会需要用到。

1
2
3
4
5
6
7
8
9
10
11
12
13
# Construct PyTorch dataset
BATCH_SIZE = 4

X_train_tensor = torch.from_numpy(X_train)
y_train_tensor = torch.from_numpy(y_train).type(torch.LongTensor)

print(X_train_tensor.shape, y_train_tensor.shape)

train_dataset = torch.utils.data.TensorDataset(X_train_tensor, y_train_tensor)
train_loader = torch.utils.data.DataLoader(train_dataset,
batch_size=BATCH_SIZE,
shuffle=True,
num_workers=2)

构建模型

下面定义一个简单的 ConvNet 类,里面实现了一个两层的 CNN。在 __init__() 里定义一些层并实现了 forward() 前向传播。因为 PyTorch 是动态图计算,因此不需要手动定义反向传播的计算过程。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
class ConvNet(nn.Module):
def __init__(self):
super(ConvNet, self).__init__()

self.conv1 = nn.Conv1d(in_channels=2,
out_channels=16,
kernel_size=3,
padding=1)
self.conv2 = nn.Conv1d(in_channels=16,
out_channels=32,
kernel_size=3,
padding=1)
self.pool = nn.MaxPool1d(2)
self.fc1 = nn.Linear(32*256, 10)

def forward(self, x):
x = self.pool(F.relu(self.conv1(x)))
x = self.pool(F.relu(self.conv2(x)))
x = x.view(-1, 32*256)
output = self.fc1(x)
return output

之后可以用 torch-summarytensorwatch 输出模型每一层的参数:

1
2
3
4
model =  ConvNet()

from torchsummary import summary
summary(model, input_size=(2, 1024))

可以得到如下的输出:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
----------------------------------------------------------------
Layer (type) Output Shape Param #
================================================================
Conv1d-1 [-1, 16, 1024] 112
MaxPool1d-2 [-1, 16, 512] 0
Conv1d-3 [-1, 32, 512] 1,568
MaxPool1d-4 [-1, 32, 256] 0
Linear-5 [-1, 10] 81,930
================================================================
Total params: 83,610
Trainable params: 83,610
Non-trainable params: 0
----------------------------------------------------------------
Input size (MB): 0.01
Forward/backward pass size (MB): 0.38
Params size (MB): 0.32
Estimated Total Size (MB): 0.70
----------------------------------------------------------------

因为是多分类任务,所以使用了 torch.nn.CrossEntropyLoss(),需要注意的是其中包含了 Softmax 因此我们不需要在网络中手动实现。优化器选择了 torch.optim.Adam,并设置学习速率为常见的 0.001

1
2
3
4
# Optimizer and criterion

criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

训练

最后我们就能实现模型训练过程,最外面是循环 EPOCHS 次,每个 epoch 中定义累积的 running_loss,前面定义的 train_loader 一次能够给出 BATCH_SIZE 的训练样本,在内循环中就能够得到训练数据并进行前向计算、计算 loss、计算梯度、优化器迭代更新整个模型的参数。

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
# Training

EPOCHS = 1

for epoch in range(EPOCHS):

model.train()
running_loss = 0.0

for batch_idx, data in enumerate(train_loader):
# get inputs and labels
inputs, labels = data

# zero the parameter gradients
optimizer.zero_grad()

# forward + backward + optimize
outputs = model(inputs)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()

# print statistics
running_loss += loss.item()

if (batch_idx+1) % 1000 == 0:
print('Epoch [{}/{}], Step [{}/{}], Loss: {:.4f}'
.format(epoch+1, EPOCHS, batch_idx+1, len(train_loader), loss.item()))

输出:

1
2
3
4
5
6
Epoch [1/1], Step [1000/6750], Loss: 2.1528
Epoch [1/1], Step [2000/6750], Loss: 1.2026
Epoch [1/1], Step [3000/6750], Loss: 1.7215
Epoch [1/1], Step [4000/6750], Loss: 1.8674
Epoch [1/1], Step [5000/6750], Loss: 1.6417
Epoch [1/1], Step [6000/6750], Loss: 1.1279

GPU 训练

如果要用 GPU 训练,需要指定 device,并且将模型以及每个 Batch 的数据都传到 GPU device 中,在上面的基础上改动或添加以下的代码:

1
2
3
4
5
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

model.to(device)

inputs, labels = data[0].to(device), data[1].to(device)

PyTorch Recipes 里给出了更多常用的例子,包括 保存 / 读取模型使用 Tensorboard 等,可以进一步参考。

参考链接