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-summary 或 tensorwatch 输出模型每一层的参数：

1
2
3
4
model =  ConvNet()

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

可以得到如下的输出：

----------------------------------------------------------------
        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 等，可以进一步参考。

PyTorch 简单分类模型示例

导入库

数据集加载

构建模型

训练

GPU 训练

参考链接

相关文章：