Intel oneDNN 深度学习库中的推理与训练编程模型解析

Intel oneDNN 深度学习库中的推理与训练编程模型解析

oneDNN 项目地址: https://gitcode.com/gh_mirrors/mk/mkl-dnn

引言

Intel oneDNN（原MKL-DNN）是一个针对深度学习应用优化的高性能计算库，专门用于加速神经网络的前向传播和反向传播计算。本文将深入解析oneDNN中关于推理（Inference）和训练（Training）的关键编程模型和技术细节。

传播类型（Propagation Kinds）

oneDNN通过不同的传播类型来区分计算场景：

推理场景

• 前向推理（forward_inference）：仅执行模型预测，不保存中间计算结果
• 特点：内存占用少，执行效率高

训练场景

1. 前向训练（forward_training）：执行预测并保存必要中间结果
2. 反向传播分为两种：
   • 反向数据传播（backward_data）：计算关于输入的梯度
   • 反向权重传播（backward_weights）：计算关于权重的梯度

训练与推理的前向传播差异

虽然数学计算相同，但实现上有重要区别：

| 特性 | 训练模式 | 推理模式 | |------|---------|---------| | 中间结果保存 | 保留必要中间结果 | 不保留 | | 内存占用 | 较高 | 较低 | | 执行速度 | 相对较慢 | 更快 | | 典型用例 | 模型训练 | 模型部署 |

以最大池化（max pooling）为例：

• 训练时需要保存最大值位置信息用于反向传播
• 推理时只需计算结果，无需保存位置信息

推理优化技巧

1. 正确使用传播类型：始终使用forward_inference
2. 内存复用：尽可能使用原地操作（in-place）
3. 持久化对象：创建一次primitive并重复使用
4. 算子融合：利用post-ops属性融合多个操作

训练特殊考量

1. 传播类型选择：

   • 前向：forward_training
   • 反向：根据需求选择backward_data或backward_weights

2. 内存格式处理：

   • 梯度张量使用format_tag::any
   • 原始数据张量需明确定义内存格式

3. 工作空间（Workspace）：

   • 某些操作需要在训练时保存中间信息
   • 工作空间在前向传播时填充，反向传播时使用

工作空间（Workspace）详解

工作空间是oneDNN中一个特殊概念，用于在训练时保存必要中间信息：

// 前向传播
auto fwd_pd = ...::primitive_desc();  // 创建primitive描述符
auto workspace_md = fwd_pd.workspace_desc();  // 获取工作空间描述
memory workspace(workspace_md, engine);  // 创建工作空间内存

fwd_primitive.execute(stream, {
        ...,
        {DNNL_ARG_WORKSPACE, workspace}  // 作为输出
        });

// 反向传播
bwd_primitive.execute(stream, {
        ...,
        {DNNL_ARG_WORKSPACE, workspace}  // 作为输入
        });

注意事项：

• 即使工作空间大小为0，创建空内存也是安全的
• 不要混淆工作空间和临时暂存空间（scratchpad）

数据类型选择

oneDNN支持多种数据类型用于训练：

• BF16（16位Brain Float）：内存占用减半，性能更高
• FP32（32位浮点）：传统选择，精度更高

最佳实践总结

1. 推理优化：

   • 使用正确传播类型
   • 最大化内存复用
   • 利用算子融合

2. 训练优化：

   • 正确处理工作空间
   • 注意内存格式一致性
   • 重用primitive对象

通过合理应用这些技术，可以充分发挥oneDNN在深度学习应用中的性能优势。

oneDNN 项目地址: https://gitcode.com/gh_mirrors/mk/mkl-dnn