找到懂机器学习系统开发的人太难了。现今的学生都一心学习机器学习算法,很多学生对于底层的运作原理理解得很浅。而当他们在实际中应用机器学习技术时才意识到系统的重要性,那时想去学习,却没有了充沛的学习时间。
导言
为了支持在不同应用中高效开发机器学习算法,人们设计和实现了机器学习框架(如TensorFlow、PyTorch、MindSpore等)。广义来说,这些框架实现了以下共性的设计目标:
神经网络编程: 深度学习的巨大成功使得神经网络成为了许多机器学习应用的核心。根据应用的需求,人们需要定制不同的神经网络,如卷积神经网络(Convolutional Neural Networks)和自注意力神经网络(Self-Attention Neural Networks)等。这些神经网络需要一个共同的系统软件进行开发、训练和部署。
自动微分: 训练神经网络会具有模型参数。这些参数需要通过持续计算梯度(Gradients)迭代改进。梯度的计算往往需要结合训练数据、数据标注和损失函数(Loss Function)。考虑到大多数开发人员并不具备手工计算梯度的知识,机器学习框架需要根据开发人员给出的神经网络程序,全自动地计算梯度。这一过程被称之为自动微分。
数据管理和处理: 机器学习的核心是数据。这些数据包括训练、验证、测试数据集和模型参数。因此,需要系统本身支持数据读取、存储和预处理(例如数据增强和数据清洗)。
模型训练和部署: 为了让机器学习模型达到最佳的性能,需要使用优化方法(例如Mini-Batch SGD)来通过多步迭代反复计算梯度,这一过程称之为训练。训练完成后,需要将训练好的模型部署到推理设备。
硬件加速器: 神经网络的相关计算往往通过矩阵计算实现。这一类计算可以被硬件加速器(例如,通用图形处理器-GPU)加速。因此,机器学习系统需要高效利用多种硬件加速器。
分布式执行: 随着训练数据量和神经网络参数量的上升,机器学习系统的内存用量远远超过了单个机器可以提供的内存。因此,机器学习框架需要天然具备分布式执行的能力。
一个完整的机器学习框架一般具有如图所示的基本架构
编程接口: 考虑到机器学习开发人员背景的多样性,机器学习框架首先需要提供以高层次编程语言(如Python)为主的编程接口。同时,机器学习框架为了优化运行性能,需要支持以低层次编程语言(如C和C++)为主的系统实现,从而实现操作系统(如线程管理和网络通讯等)和各类型硬件加速器的高效使用。
计算图: 利用不同编程接口实现的机器学习程序需要共享一个运行后端。实现这一后端的关键技术是计算图技术。计算图定义了用户的机器学习程序,其包含大量表达计算操作的算子节点(Operator Node),以及表达算子之间计算依赖的边(Edge)。
编译器前端: 机器学习框架往往具有AI编译器来构建计算图,并将计算图转换为硬件可以执行的程序。这个编译器首先会利用一系列编译器前端技术实现对程序的分析和优化。编译器前端的关键功能包括实现中间表示、自动微分、类型推导和静态分析等。
编译器后端和运行时: 完成计算图的分析和优化后,机器学习框架进一步利用编译器后端和运行时实现针对不同底层硬件的优化。常见的优化技术包括分析硬件的L2/L3缓存大小和指令流水线长度,优化算子的选择或者调度顺序。
异构处理器: 机器学习应用的执行由中央处理器(Central Processing Unit,CPU)和硬件加速器(如英伟达GPU、华为Ascend和谷歌TPU)共同完成。其中,非矩阵操作(如复杂的数据预处理和计算图的调度执行)由中央处理器完成。矩阵操作和部分频繁使用的机器学习算子(如Transformer算子和Convolution算子)由硬件加速器完成。
数据处理: 机器学习应用需要对原始数据进行复杂预处理,同时也需要管理大量的训练数据集、验证数据集和测试数据集。这一系列以数据为核心的操作由数据处理模块(例如TensorFlow的tf.data和PyTorch的DataLoader)完成。
模型部署: 在完成模型训练后,机器学习框架下一个需要支持的关键功能是模型部署。为了确保模型可以在内存有限的硬件上执行,会使用模型转换、量化、蒸馏等模型压缩技术。同时,也需要实现针对推理硬件平台(例如英伟达Orin)的模型算子优化。最后,为了保证模型的安全(如拒绝未经授权的用户读取),还会对模型进行混淆设计。
分布式训练: 机器学习模型的训练往往需要分布式的计算节点并行完成。其中,常见的并行训练方法包括数据并行、模型并行、混合并行和流水线并行。这些并行训练方法通常由远端程序调用(Remote Procedure Call, RPC)、集合通信(Collective Communication)或者参数服务器(Parameter Server)实现。