Lab5-简单神经网络训练与加速

实验简介

1
2
3
4
5
6
深度学习（Deep Learning）是机器学习的分支，是一种以人工神经网络为架构，对数据进行表征学习的算法。深度学习
能够取得如此卓越的成就，除了优越的算法、充足的数据，更离不开强劲的算力。近年来，深度学习相关的基础设施逐渐
成熟，从网络设计时的训练、优化，到落地的推理加速，都有非常优秀的解决方案。其中，对于算力的需求最大的部分之
一是网络的训练过程，它也因此成为 HPC 领域经常研究的话题。

本次实验我们将完成 LeNet-5 的训练，并尝试编写自定义算子。

实验环境

  • 在 H248 节点上用 conda activate torch 导入环境
  • 用 sbatch 脚本任务运行程序
1
2
3
4
5
6
7
8
9
#!/bin/bash
#SBATCH -o out.txt
#SBATCH -N 1
#SBATCH -n 1
#SBATCH -p 2080Ti
#SBATCH --cpus-per-task=24
#SBATCH --gpus=1

CUDA_VISIBLE_DEVICES=1 python LeNet-5.py

LeNet-5 训练

  • 数据准备
1
2
3
4
5
6
7
# 加载MNIST数据集
train_dataset = torchvision.datasets.MNIST(root = '../Lab5/data/', train = True, transform = transforms.ToTensor(), download = True)
test_dataset = torchvision.datasets.MNIST(root = '../Lab5/data/', train = False, transform = transforms.ToTensor())

# 定义数据加载器
train_loader = torch.utils.data.DataLoader(dataset = train_dataset, batch_size = 64, shuffle = True)
test_loader = torch.utils.data.DataLoader(dataset = test_dataset, batch_size = 64, shuffle = False)

  • 网络结构

    LeNet-5 神经网络原版结构除去输入层外共有 7 层，分别为 2 个卷积层、2 个池化层和 3 个全连接层，激活函数这里选取的是 GELU ，网络结构的构建参考了网上的代码。

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
class LeNet5(nn.Module):
    def __init__(self):
        super(LeNet5, self).__init__()
        # 第一层卷积，输入通道为 1，输出通道为 6，卷积核大小为 5 * 5，步距为 1
        self.conv1 = nn.Conv2d(in_channels = 1, out_channels = 6, kernel_size = 5, stride = 1)
        # 第一层池化层，卷积核为 2 * 2，步距为 2
        self.pool1 = nn.MaxPool2d(kernel_size = 2, stride = 2)
        # 第二层卷积，输入通道为 6，输出通道为 16，卷积核大小为 5 * 5，步距为 1
        self.conv2 = nn.Conv2d(in_channels = 6, out_channels = 16, kernel_size = 5, stride = 1)
        # 第二层池化层，卷积核为 2 * 2，步距为 2
        self.pool2 = nn.MaxPool2d(kernel_size = 2, stride = 2)
        # 第一层全连接层，维度由 16 * 4 * 4 => 120
        self.fc1 = nn.Linear(in_features = 16 * 4 * 4, out_features = 120)
        # 第二层全连接层，维度由 120 => 84
        self.fc2 = nn.Linear(in_features = 120, out_features = 84)
        # 第三层全连接层，维度由 84 => 10
        self.fc3 = nn.Linear(in_features = 84, out_features = 10)
        
    def forward(self, x):
        # 将数据送入第一个卷积层，再送入第一个池化层
        x = self.pool1(F.gelu(self.conv1(x)))
        # 将数据送入第二个卷积层，再送入第二个池化层
        x = self.pool2(F.gelu(self.conv2(x)))
        # 将池化层后的数据进行 reshape
        x = x.view(-1, 16 * 4 * 4)
        # 将数据送入第一个全连接层
        x = F.gelu(self.fc1(x))
        # 将数据送入第二个全连接层
        x = F.gelu(self.fc2(x))
        # 将数据送入第三个全连接层得到输出
        x = self.fc3(x)
        return x
  • 定义模型、损失函数和优化器
1
2
3
4
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = LeNet5().to(device)
criterion = nn.CrossEntropyLoss().to(device)
optimizer = torch.optim.Adam(model.parameters())
  • 训练过程
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
for i, (x_train, y_train) in enumerate(train_loader):
    # 将数据转移到 GPU 上
    x_train = x_train.to(device)
    y_train = y_train.to(device)

    # 执行 forward，通过将 x_train 传递给模型来计算 y_pred
    y_pred = model(x_train)

    # 计算损失
    loss = criterion(y_pred, y_train)
    
    # 清零梯度，执行 backward，并更新权重
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
  • 测试过程
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
model.eval()
with torch.no_grad():
    correct = 0
    total = 0
    for x_test, y_test in test_loader:
        # 将数据转移到 GPU 上
        x_test = x_test.to(device)
        y_test = y_test.to(device)

        # 执行 forward，通过将 x_test 传递给模型来计算 y_pred
        y_pred = model(x_test)

        # 得到预测结果
        _, predicted = torch.max(y_pred.data, 1)

        # 统计 total 和 correct
        total += y_test.size(0)
        correct += (predicted == y_test).sum().item()

    print('Test Accuracy: {:.2f}%'.format(100 * correct / total))
  • 测试结果

  • GPU 占用率

    sbatch 脚本任务在跑的时候 ssh 到节点上用 nvidia-smi 命令查看

  • TensorBoard

    使用 pip install tensorboard 命令安装 TensorBoard 。

1
2
3
4
5
6
# 导入 TensorBoard
from torch.utils.tensorboard import SummaryWriter
# 创建 TensorBoard 写入器
writer = SummaryWriter('../Lab5/Tensorboard')
# 关闭 TensorBoard 写入器
writer.close()
1
2
3
4
# 记录损失
writer.add_scalar('Loss', average_loss / 100, epoch * len(train_loader) + i)
# 记录正确率
writer.add_scalar('Accuracy', 100 * correct / total, epoch)

在集群上跑完一次任务后会产生记录文件，使用 tensorboard --logdir=./Lab5/Tensorboard 命令后默认会在 localhost:6006 上打开，本地新建一个终端用 ssh -L 8080:localhost:6006 zsh@clusters.zju.edu.cn -p 80 转发端口，然后在本地浏览器上打开 localhost:8080 就可以看到图像。

编写 GELU 算子

  • forward 函数

    forward 函数接受输入并返回输出，这里用的是 GELU 的近似版本 \[ \rm{GELU}(x) = 0.5 * x * (1 + \tanh(\sqrt{\frac{2}{\pi}} * (x + 0.044715 * x^3))) \]

1
2
3
def forward(ctx, input):
    ctx.save_for_backward(input)
    return 0.5 * input * (1 + torch.tanh(math.sqrt(2 / math.pi) * (input + 0.044715 * torch.pow(input, 3))))

  • backward 函数

    backward 函数接受输入和梯度，并返回相对于输入的梯度，这里使用 WolframAlpha 对 GELU 函数进行了微分，得到以下结果，根据链式法则将 grad_output 与之相乘就可以得到 grad_output。

\[ \begin{gather*} \rm{GELU}'(x) = 0.5 * \tanh(0.0356774 * x^3 + 0.797885 * x)\\ + \rm(0.0535161 * x^3 + 0.398942 * x) * sech^2(0.0356774 * x^3 + 0.797885 * x) + 0.5 \end{gather*} \]

1
2
3
4
5
6
7
def backward(ctx, grad_output):
    input, = ctx.saved_tensors
    part1 = 0.5 * torch.tanh(0.0356774 * torch.pow(input, 3) + 0.797885 * input)
    part2 = 0.0535161 * torch.pow(input, 3) + 0.398942 * input
    part3 = 1.0 / (torch.cosh(0.0356774 * torch.pow(input, 3) + 0.797885 * input))
    grad_input = grad_output * (part1 + part2 * torch.pow(part3, 2) + 0.5)
    return grad_input

  • 正确性验证

    使用以下代码来验证正确性

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
A = torch.randn(100)
B = A.clone()
A.requires_grad = True
B.requires_grad = True
c = torch.randn(100)
a = F.gelu(A)
b = my_gelu(B)
loss1 = F.mse_loss(a, c)
loss2 = F.mse_loss(b, c) # loss1 should equal to loss2
loss1.backward()
loss2.backward()
gradA = A.grad
gradB = B.grad
err = F.mse_loss(gradA, gradB) # err should equal to 0
print(gradA)
print('\n')
print(gradB)
print('\n')
print(err)

使用 C++ 算子训练 LeNet-5

文件夹下共有三个文件：setup.py、C_my_gelu.cpp、LeNet-5.py

  • setup.py

    用 C_my_gelu.cpp 来创建 C_my_gelu_cpp 模块

1
2
3
4
5
6
from setuptools import setup, Extension
from torch.utils import cpp_extension

setup(name = 'C_my_gelu_cpp',
      ext_modules = [cpp_extension.CppExtension('C_my_gelu_cpp', ['C_my_gelu.cpp'])],
      cmdclass = {'build_ext': cpp_extension.BuildExtension})

  • C_my_gelu.cpp

    首先需要引入 <torch/extension.h> 头文件，然后构建 forward 和 backward 函数，这里与 python 自定义算子的实现是一致的,并且要完成 pybind11 绑定

1
2
3
4
// 正向传播
torch::Tensor gelu_forward(const torch::Tensor& input) {
    return 0.5 * input * (1 + torch::tanh(sqrt(2.0 / pi) * (input + 0.044715 * torch::pow(input, 3))));
}
1
2
3
4
5
6
7
8
// 反向传播
torch::Tensor gelu_backward(const torch::Tensor& input, const torch::Tensor& grad_output) {
    torch::Tensor part1 = 0.5 * torch::tanh(0.0356774 * torch::pow(input, 3) + 0.797885 * input);
    torch::Tensor part2 = 0.0535161 * torch::pow(input, 3) + 0.398942 * input;
    torch::Tensor part3 = 1.0 / torch::cosh(0.0356774 * torch::pow(input, 3) + 0.797885 * input);
    torch::Tensor grad_input = grad_output * (part1 + part2 * torch::pow(part3, 2) + 0.5);
    return grad_input;
}
1
2
3
4
5
// pybind11 绑定
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m){
    m.def("forward", &gelu_forward, "gelu forward");
    m.def("backward", &gelu_backward, "gelu backward");
}

  • 使用 python setup.py install 命令构建 C_my_gelu_cpp 模块

  • LeNet-5.py

    前面 setup.py 构建了名为 C_my_gelu_cpp 的模块，LeNet-5.py 中通过导入就可以使用其中的函数。然后构建 C_my_gelu 新模块并把原来 LeNet-5 神经网络训练过程中的 F.gelu 激活函数替换成 my_gelu 就完成了自定义算子训练

1
2
import C_my_gelu_cpp
from torch.autograd import Function
1
2
3
4
5
6
7
8
9
10
11
12
class C_my_gelu(Function):
    @staticmethod
    def forward(ctx, input):
        ctx.save_for_backward(input)
        return C_my_gelu_cpp.forward(input)

    @staticmethod
    def backward(ctx, grad_output):
        input, = ctx.saved_tensors
        return C_my_gelu_cpp.backward(input, grad_output)

my_gelu = C_my_gelu.apply

  • 正确性验证

    使用与上面相同的代码来验证正确性

  • 测试结果