编辑
2024-07-06
学习记录
0
请注意,本文编写于 338 天前,最后修改于 338 天前,其中某些信息可能已经过时。

目录

高性能计算
AI框架相关知识
1.MLIR概念
2.ONNX概念
3.TVM的整体结构,如何用TVM进行开发?
4.为什么要进行推理优化?直接用tensorflow或者pytorch的推理接口不行吗?
5.模型推理优化的常用方法有哪些?
6.TensorRT如何进行自定义算子开发
7.TensorRT对模型实现了哪些推理优化
8.算子融合为什么能加速推理,优化了哪一部分?TensorRT用到了哪些算子融合?算子融合在推理框架中是如何实现的?
9.模型量化的加速原理,模型量化带来的精度损失如何解决?
10.ONNX Runtime支持在多种硬件上进行推理,说明具体的实现机制。
11.总结一下TensorRT,ONNX Runtime等推理框架的组成架构,如果我们公司自己要为硬件开发一套推理框架,应该重点关注哪些部分?
12.各种推理框架都有何优劣势?它们的性能怎么样?
13.分布式训练中有哪些并行模式?每种模式需要做什么,有什么优缺点?
14.分布式训练中我们重点需要处理的问题有哪些?
15.MPI(消息传递推理)可以应用于AI框架的哪些方面?
16.反向传播的原理,具体实现的源码有了解过吗?
17.TensorFlow和Pytorch都用过吗?他们的设计思路有何不同?有何优劣?如何添加自定义算子?

高性能计算

AI框架相关知识

1.MLIR概念

MLIR是一个多级中间表示框架,旨在解决当前深度学习领域中存在的IR(Intermediate Representation,中间表示)碎片化问题。它由LLVM团队开发和维护,强调工具链的可重用性和可扩展性。MLIR试图通过引入统一的表示形式——Dialect(方言),来克服不同IR之间转换难度大、优化Pass难以共享等问题。

MLIR它不仅仅局限于深度学习或某个特定的框架。MLIR旨在提供一个多级抽象的表示,支持从高层次算法到低层硬件指令的各种编程范式和优化技术。MLIR既可以用于深度学习模型的表示和优化,也可以用于传统软件编译领域,甚至是硬件设计的高级表示。

2.ONNX概念

ONNX(Open Neural Network Exchange)是一种开放的IR,主要用于不同深度学习框架之间的模型转换和互操作性。ONNX通过定义一套标准化的算子和模型格式来实现这一点。

无论你使用何种训练框架训练模型(比如TensorFlow/Pytorch/OneFlow/Paddle),在训练完毕后你都可以将这些框架的模型统一转换为ONNX这种统一的格式进行存储。

然而,ONNX主要关注模型的表示和交换,而不涉及模型的优化过程。

3.TVM的整体结构,如何用TVM进行开发?

image.png

TVM是一个端到端的深度学习编译框架,整体结构可以分为以下几个关键组成部分:

  1. 前端:TVM支持多种深度学习框架的模型作为输入,比如TensorFlow、PyTorch、MXNet、Keras等。通过这些前端接口,TVM可以读取不同框架定义的模型,并将其转换成中间表示(IR)。
  2. 中间表示(IR):TVM使用两级IR,即Relay和TIR(Tensor IR)。Relay是一种高级IR,用于表示高级神经网络算法;而TIR是一种低级IR,用于表示更接近于硬件的操作和优化。
  3. 自动调度(AutoTVM/AutoScheduler):为了在特定硬件上获得最佳性能,TVM提供了自动调度工具,如AutoTVM和更现代的AutoScheduler,它们可以自动优化模型的计算图和内核实现。
  4. 运行时:TVM提供了一个轻量级的运行时,支持模型在目标硬件上的部署和执行。这包括对多种设备的支持,如CPU、GPU、FPGA等。
  5. 编译流程:TVM的编译流程包括模型的加载、优化(例如算子融合、内存优化)、自动调度、代码生成等步骤,最终生成可以在目标硬件上运行的机器码。

4.为什么要进行推理优化?直接用tensorflow或者pytorch的推理接口不行吗?

首先推理优化的原因主要是为了提高模型在实际应用中的性能,如快速响应。

其次,随着深度学习的广泛应用和以Transformer大模型为基座的深度算法的普及,tensorflow或者pytorch的推理性能远远跟不上了。

为了提高模型在实际应用中的性能,解决其在推理时面临的高计算复杂度、大内存需求、并行性限制等问题,常用一些办法总结如下:

  1. 模型压缩技术:通过知识蒸馏、权重剪枝、量化等方法,减小模型体积和计算需求。这些技术通过精简模型结构或参数,减少不必要的计算负担,从而加快推理速度,降低存储和运行时内存需求。知识蒸馏特别值得注意,它通过将大模型的"知识"转移到小模型上,既保持了模型性能,又实现了模型大小的显著减少。
  2. 高效变体的开发:针对自注意力机制的高计算复杂度,开发了如Transformers的变体,通过修改自注意力机制来降低计算复杂度和内存需求。例如,使用稀疏注意力模式、局部注意力或低秩近似等技术,有效地减少了处理长序列时的计算负担。
  3. 硬件加速和专用推理引擎:利用TPU、GPU等专门设计的硬件加速器和优化的推理引擎(如TensorRT、ONNX Runtime),针对特定硬件平台的特性进行底层优化。这种优化可以显著提高推理速度,同时降低功耗,特别适合需要实时处理的应用场景。
  4. 动态量化和混合精度推理:通过在模型的不同部分使用不同的数据精度(如FP32、FP16、INT8)来平衡推理速度和模型精度之间的关系。动态量化特别适用于在运行时根据需求调整精度,而混合精度推理则可以在保持模型性能的同时加速模型的推理过程。

5.模型推理优化的常用方法有哪些?

推理性能优化的常用方法主要包括两大类,一类是模型的压缩技术,比如模型的剪枝、量化、蒸馏等,在较大的预训练模型下被广泛使用;另外是推理加速的技术,在CPU、GPU加速上分别有一些方法。

1. 网络剪枝

网络剪枝是从大型网络中筛选出不重要的神经元以及权重,将它们从网络中删除,同时尽可能地保留网络的性能。

2. 量化技术

量化技术是把高精度表示的网络权重和激活值,用低精度来近似表示,实现网络的轻量化。优势如下:

网络存储 :每个层权重量化后,32位的比特就可以压缩到8比特,就是浮点型到整形的量化,整个模型占的空间就会变小;

激活值 : 通过使用较少位的数值表示,在处理同样数据时需要读/写的内容就更短,内存带宽的压力就变得更小;

计算时间 :单位时间内处理定点运算指令就会比浮点运算的指令多。

3. 模型蒸馏

蒸馏是一种模型压缩常见方法,将复杂、学习能力强的网络学到的特征,表示“知识”蒸馏出来,传递给参数量小、学习能力弱的网络。

4. Caching

CPU频率远快于主存访问速度,在处理器时钟周期内,CPU常常需要等待主存,浪费计算资源。为了缓解CPU和内存之间速度的不匹配问题,增加了CPU cache 来解决。

cache 利用局部性原理来提高缓存命中率:

  1. 时间局部性:如果某个数据被访问,那么在不久的将来它很可能被再次访问;
  2. 空间局部性:如果某个数据被访问,那么与它相邻的数据很快也可能被访问。

5. 多算子融合

image.png 算子融合是GPU上一个很重要的推理加速的一个优化手段,尤其是针对NLP这样的大模型,会带来比较显著的效果的提升。对于GPU异构编程,每一次op操作都会有一个内核的调用和多次的显存的读取;对于小op来说启动GPU kernel的时间会大于GPU计算时间,显存的读取开销也很大;op数目太多的话,效率会变低;所以将算子合并,可以有效地提高计算的性能。

6. 计算图优化

每个计算图中都包含许多计算节,图优化的目标很简单,就是简化计算图中计算节点的计算量。常用的方式分为以下几种:

  1. 减少节点的数量
  2. 用高效替换低效的节点
  3. 用高效子图替换低效子图
  4. 用并行化分支代替单分支

6.TensorRT如何进行自定义算子开发

首先,明确为什么需要自定义算子,了解TensorRT中自定义算子的基本概念,包括插件(Plugin)和插件工厂(Plugin Factory)。

其次,开始实现自定义算子,主要分为三步,定义插件类,实现算子逻辑和注册插件,下面分别展开说说。

  1. 定义插件类:继承nvinfer1::IPluginV2接口,实现其虚函数,包括算子计算(enqueue)、输出维度(getOutputDimensions)、数据类型和格式配置(configureWithFormat)、初始化与清理(initialize 和 terminate)、序列化与反序列化(serialize 和 deserialize)等。
  2. 实现算子逻辑:核心实现:在enqueue函数中实现算子的具体计算逻辑,可能需要使用CUDA等技术在GPU上进行并行计算以提高效率。
  3. 注册插件:通过实现nvinfer1::IPluginCreator接口并注册插件,使得TensorRT能够识别和使用自定义算子。

再次,是构建使用自定义算子代码

  1. 创建插件实例:在网络构建过程中,使用插件工厂或直接调用插件类的构造函数创建自定义算子实例,并将其嵌入到模型中。
  2. 序列化和反序列化:确保自定义算子能够被正确地序列化和反序列化,以便模型的保存和加载。

最后,测试和验证

验证正确性和性能:通过与预期结果或其他框架的对比,验证自定义算子的实现是否正确、性能是否达标,并确保其在不同条件下的鲁棒性。

7.TensorRT对模型实现了哪些推理优化

  1. 常量折叠(Constant Folding)

解释:常量折叠是一种编译时优化技术,它预计算图中那些在推理前就能确定结果的表达式。这意味着网络中的任何常量操作,如常量之间的算术运算,都会在模型编译期间被提前计算并简化,从而减少运行时的计算负担。

好处:通过减少不必要的运行时计算,可以显著提高模型的推理速度。

  1. 算子融合

解释:算子融合是将多个操作合并为一个复合操作的过程,这可以减少内存访问次数并降低推理延迟。例如,卷积、批量归一化(Batch Normalization)、激活函数等可以融合成一个单一的高效操作。

好处:减少了内存访问和计算步骤,提高了数据吞吐率和运算效率。

  1. 量化

解释:量化是指将模型中的权重和激活从浮点数转换为低精度的表示形式,如从32位的浮点数(FP32)转换为8位整数(INT8)。TensorRT提供了量化校准工具,以最小化量化带来的精度损失。

好处:量化可以显著减少模型的大小和推理时间,同时也减少了内存带宽的需求,使得模型更适合在资源受限的设备上运行。

  1. 层自动调整

解释:TensorRT会针对特定的硬件平台自动选择最优的算法来执行各种操作。这是通过在不同的算法实现之间运行基准测试来完成的,以确保选择的实现能够提供最佳性能。

好处:确保模型在特定硬件上达到最优的运行效率。

  1. 多流执行(Multi-Stream Execution)

解释:TensorRT支持利用GPU的并行处理能力,通过多流执行来同时处理多个推理请求,从而提高吞吐量。

好处:在处理大量并发请求时,能够有效提升GPU利用率和总体吞吐量。

8.算子融合为什么能加速推理,优化了哪一部分?TensorRT用到了哪些算子融合?算子融合在推理框架中是如何实现的?

首先,算子融合通过将多个操作合并为一个单一复合操作来减少模型中的层次。这种优化减少了内存访问次数和数据传输量,因为它降低了中间结果的读写需求。同时,减少了CPU和GPU之间的同步点,提高了计算效率。

其次,TensorRT利用算子融合技术,例如将卷积、批量归一化(Batch Normalization)、激活函数(如ReLU)融合为一个单一操作。这种融合不仅减少了计算步骤,还减少了中间数据的存储需求。

在推理框架中,算子融合通常在图优化阶段进行。框架会分析计算图,识别可以融合的操作序列。然后,通过生成一个包含所有融合操作逻辑的新内核来实现融合。这个过程可能涉及到自动生成代码或者使用预先定义的高效内核。

然而,实现算子融合需要考虑操作之间的依赖关系和数据流动,确保融合后的操作不会引入错误。此外,还需要平衡融合的程度和实现的复杂性,以达到最佳性能。

9.模型量化的加速原理,模型量化带来的精度损失如何解决?

首先,模型量化是指将模型中的权重和激活函数的数据类型从浮点数(如FP32)转换为低精度的格式(如INT8或FP16)。

这种转换减少了模型的内存占用,降低了计算所需的内存带宽,并且在一些硬件上可以利用专门的低精度计算单元,从而加速模型的推理过程。

量化可以显著减少模型大小,提高数据加载和处理速度。例如,在支持INT8指令集的硬件上,使用INT8量化的模型相比于FP32模型,理论上可以提高4倍的内存带宽效率和计算速度。

量化带来的精度损失及解决方案:

精度损失原因:量化过程中,由于将浮点数映射到有限的整数范围,会引入量化误差,导致模型的精度下降。

解决策略:

  1. 量化校准:通过在量化前后对模型进行校准,选择最优的量化参数(如量化比例和零点),以最小化量化误差。常用的校准方法包括最小最大值校准、百分位校准等。
  2. 量化感知训练:在模型训练过程中模拟量化的效果,让模型“适应”量化带来的误差。这种方法可以在训练阶段就考虑到量化误差,进而学习到更鲁棒的权重。

模型量化是一种有效的推理加速技术,虽然可能会带来一定的精度损失,但通过量化校准和量化感知训练等策略,可以显著减轻甚至克服这一问题。

10.ONNX Runtime支持在多种硬件上进行推理,说明具体的实现机制。

首先,ONNX Runtime是一个用于优化和运行机器学习模型的性能引擎。它支持使用ONNX格式的模型,这是一个开放格式,用于表示机器学习模型,使得不同的AI框架训练的模型能够在不同的平台和设备上运行。

其次,多硬件支持的实现机制实现原理是,ONNX Runtime通过提供一系列的“执行提供程序”(Execution Providers,EPs)来支持不同的硬件。每个执行提供程序都是为特定的硬件或计算库优化的后端,它定义了如何在该硬件上执行ONNX模型中的操作。它会根据可用的执行提供程序和硬件资源自动选择最适合当前环境的执行路径。用户也可以手动指定使用哪个执行提供程序,以便更精细地控制模型的运行方式。

例如,NVIDIA GPU的系统上,ONNX Runtime可以利用CUDA执行提供程序来加速模型的推理。此外,对于需要进一步优化的场景,可以使用TensorRT执行提供程序,它利用NVIDIA TensorRT进行图优化和内核融合,以实现更高效的推理。

总之,ONNX Runtime通过灵活的执行提供程序机制,有效地支持了多种硬件平台,这种跨平台的能力大大降低了模型部署的复杂性,并为开发者提供了更多的灵活性和选择。

11.总结一下TensorRT,ONNX Runtime等推理框架的组成架构,如果我们公司自己要为硬件开发一套推理框架,应该重点关注哪些部分?

从TensorRT,ONNX Runtime推理框架来看,推理框架的核心组成部分应该分为以下几个部分:

  1. 模型解析器:负责将机器学习模型(如ONNX格式)转换成框架能够理解和执行的内部格式。
  2. 图优化器:通过算子融合、常数折叠等技术优化计算图,减少不必要的计算,提高执行效率。
  3. 执行引擎:负责根据优化后的计算图,在特定硬件上执行计算任务。
  4. 硬件抽象层(对应于TensorRT、ONNX Runtime中的执行提供者/后端):为不同硬件提供定制化支持,确保模型能够在多种平台上高效运行。

在为硬件开发一套推理框架时,重点要关注两个方面,兼容和利用硬件特点优化性能和图优化技术。

推理框架需要根据目标硬件的特点进行专门的优化,通过利用硬件的并行计算能力、特殊的指令集、高效的内存访问,显著提高模型的推理速度,降低能耗,突出差异化和优势。

通过算子融合、优化数据传输路径等,减少不必要的计算和内存访问,提升模型执行的效率。这对于加速模型推理、减少资源消耗至关重要。

12.各种推理框架都有何优劣势?它们的性能怎么样?

  1. TensorRT
  • 优势:专为NVIDIA GPU设计,提供高度优化的推理性能,特别适用于高吞吐量的服务器端和边缘设备场景。支持精确的算子融合、量化、动态张量等高级优化。
  • 劣势:主要限制在NVIDIA GPU上,不适用于其他类型的硬件。相对其他框架,学习曲线可能更陡峭。
  • 性能:在NVIDIA GPU上,TensorRT通常能提供最佳的推理速度和效率。
  1. ONNX Runtime
  • 优势:支持多种硬件平台,包括CPU、GPU和FPGA等。与ONNX模型格式紧密集成,方便从不同的训练框架迁移模型。微软背书,社区活跃。
  • 劣势:虽然支持多种硬件,但在特定硬件上的优化可能不如专门的推理引擎深入。
  • 性能:提供良好的跨平台性能,但在特定硬件上可能不是最优。
  1. TensorFlow Lite
  • 优势:专为移动和嵌入式设备优化,支持模型量化和轻量级操作,减少模型大小和提升推理速度。广泛应用于Android和iOS设备。
  • 劣势:相比于服务器端的框架,可能在功能和性能上有所限制。
  • 性能:在移动设备上提供优化的推理性能,特别是在支持NEON指令集和GPU加速的设备上。
  1. Core ML
  • 优势:苹果官方支持,为iOS、macOS、watchOS和tvOS设备上的机器学习应用优化。可以直接利用苹果设备的硬件加速能力。
  • 劣势:限于苹果生态系统,不适用于非苹果平台。
  • 性能:在苹果设备上提供高效的推理性能,特别是通过使用Metal进行GPU加速时。

性能总结: 推理框架的性能受到多种因素影响,包括模型的复杂度、硬件的计算能力、框架的优化程度等。一般而言,专门为特定硬件优化的推理框架(如TensorRT、Core ML)能够提供最佳性能。然而,跨平台框架(如ONNX Runtime、TensorFlow Lite)提供了更广泛的适用性和灵活性,允许模型在多种设备上运行,虽然可能牺牲一定的性能。

13.分布式训练中有哪些并行模式?每种模式需要做什么,有什么优缺点?

分布式训练主要分为两种数据并行和模型并行。

数据并行适用于大规模训练数据集和相对较小的模型,能够提高训练速度。通常需要将训练数据集划分为多个子集,每个设备或节点负责处理一个子集,并在每个子集上独立训练模型。然后,通过梯度聚合和同步来更新模型参数。

它的缺点是需要大量的通信开销,因为设备之间需要传输梯度信息。同时,需要额外的内存来存储模型副本和梯度。

模型并行适用于大型模型或需要更高计算需求的任务,允许训练更大规模的模型。它指的是将模型分解为多个部分,在不同的设备或节点上并行处理。每个设备只负责处理模型的一部分,并与其他设备交换中间结果。

缺点是需要更复杂的编程和通信模式,以确保各个设备之间的协同工作。可能存在设备之间的通信瓶颈。

不过在实际应用中,一般是可以将数据并行和模型并行结合使用,以充分利用多个设备和节点的计算能力。

14.分布式训练中我们重点需要处理的问题有哪些?

  1. 通信开销和延迟:

在分布式环境下,非常需要关注设备之间的通信开销和传输延迟。为了减少通信开销,可以采用压缩、稀疏化和量化等技术。同时,可以利用高速网络和专用硬件来减少传输延迟。

例如,可以使用混合精度训练(Mixed Precision Training)将参数和梯度从单精度浮点数压缩为半精度浮点数,从而减少传输数据的大小。在多机之间进行数据传输时,可以采用RDMA(Remote Direct Memory Access)网络或InfiniBand网络等减少传输延迟。

  1. 容错性和可伸缩性

在分布式环境下,尤其是多机训练中,必须考虑的就是训练中的设备故障或网络异常。例如断点训练,冗余参数服务器等。

15.MPI(消息传递推理)可以应用于AI框架的哪些方面?

  1. 数据并行性:

AI模型的训练通常需要处理大规模的数据集,而数据并行性可以将数据划分为多个部分,并将其分发到不同的计算节点上并行处理。MPI提供了消息传递的能力,可以在不同计算节点之间传输数据,实现数据的并行处理。

  1. 参数同步:

在分布式训练中,不同计算节点上的模型参数需要进行同步,以保持一致性和提高训练效果。MPI可以通过消息传递的方式,在计算节点之间进行参数的同步和更新,确保所有节点上的模型参数保持一致。

  1. 集群扩展:

MPI可以应用于AI框架中的集群扩展,使得框架可以在大规模分布式环境下运行。通过MPI,可以实现节点之间的通信和数据传输,以及任务的分发和调度,从而实现在集群中进行规模化的训练和推理。

  1. 容错性:

AI框架中的分布式训练通常需要考虑设备故障和恢复的情况。MPI提供了容错功能,可以在计算节点故障后重新分配任务和数据,并进行恢复,以保证训练过程的连续性和可靠性。

可以参考Horovod在TensorFlow、PyTorch深度学习训练中的使用。

16.反向传播的原理,具体实现的源码有了解过吗?

损失函数求得的结果对权重和偏置求偏导,和原来的参数做差

深度神经网络的权重是如何精确调整的?它们就是通过 反向传播进行调整的。如果没有反向传播,深度神经网络将无法执行识别图像和解释自然语言等任务。

反向传播的目标 是为了减少损失或误差,通过调整网络的权重来实现的,使假设更像输入特征之间的真实关系。

前向传播是神经网络中的常规训练过程,通过对输入进行加权求和和激活函数运算得到输出,然后将该输出作为下一层的输入,一直传递到神经网络的末端。

反向传播过程接受模型训练过程的最终决策,然后确定这些决策中的错误。通过对比网络的输出/决策和网络的预期/期望输出来计算误差。

一旦计算出网络决策中的错误,该信息就会通过网络反向传播,并且网络参数会随之改变。用于更新网络权重的方法基于微积分,具体来说,它基于链式法则。

当神经元提供输出值时,会使用传递函数计算输出值的斜率,从而产生派生输出。进行反向传播时,特定神经元的误差根据以下公式计算 公式:

误差 = (预期输出 – 实际输出) * 神经元输出值的斜率

对输出层的神经元进行操作时,将类别值作为期望值。计算出误差后,该误差将用作隐藏层中神经元的输入,这意味着该隐藏层的误差是输出层中找到的神经元的加权误差。误差计算沿着权重网络向后传播。

计算出网络的误差后,必须更新网络中的权重。如前所述,计算误差涉及确定输出值的斜率。计算出斜率后,将进行一个称为 梯度下降 可用于调整网络中的权重。

梯度是一个斜率,其角度/陡度可以测量。斜率是通过在“运行”上绘制“y”或“上升”来计算的。在神经网络和错误率的情况下,“y”是计算的误差,而“x”是网络的参数。网络的参数与计算的误差值有关,并且随着网络权重的调整,误差会增加或减少。

“梯度下降”是更新权重以降低错误率的过程。反向传播用于预测神经网络参数与错误率之间的关系,从而建立梯度下降网络。使用梯度下降训练网络涉及通过前向传播计算权重、反向传播误差,然后更新网络的权重。

17.TensorFlow和Pytorch都用过吗?他们的设计思路有何不同?有何优劣?如何添加自定义算子?

TensorFlow像是搞数据的人研发的,Pytorch才像是搞算法人研发的。

Tensorflow 和 pytorch 相比,前者占据先机,后者则势头完全盖过前者。

TensorFlow:由Google开发,设计时考虑了分布式计算、大规模数据处理和生产环境的需求。TensorFlow使用静态计算图,需要先定义后运行,适合于大型模型和复杂的神经网络。

PyTorch:由Facebook的人工智能研究团队开发,设计理念侧重于灵活性和直观性。PyTorch采用动态计算图,允许即时修改和执行,非常适合于快速原型开发和研究工作。

TensorFlow:虽然提供了强大的功能,但其静态图机制使得调试和理解模型相对复杂。TensorFlow 2.0及后续版本引入了Eager Execution,改善了易用性。但API得兼容性也是一个大问题。

PyTorch:动态图的特性使得PyTorch在模型构建和调试方面更加直观和用户友好。PyTorch的API设计更接近Python原生,易于理解和使用。

TensorFlow:在大规模数据集和复杂模型的训练上展现出较强的性能,尤其是在GPU加速和TPU支持方面。

PyTorch:虽然在早期版本中性能略逊于TensorFlow,但近年来通过优化和社区努力,其性能已大幅提升,尤其在某些特定任务上表现出色。

TensorFlow:由于其稳定性和规模化部署能力,非常适合用于商业产品和大型企业项目。TensorFlow Serving、TensorFlow Lite和TFX等工具支持了从研发到生产的全流程。

PyTorch:以其灵活性和友好的API,更受研究人员和数据科学家的喜爱,特别是在进行快速实验和研究原型开发时。不过随着大语言模型的爆火,很多大语言模型的商用都是基于PyTorch。

在TensorFlow和PyTorch中,可以通过以下步骤来添加自定义算子:

  1. TensorFlow:

首先,是实现底层算子,可以使用C++或者CUDA,来写出你想要实现的功能,也就是自定义算子的代码。把这段代码编译成一个库文件,这样我们的TensorFlow才能认识它。

其次,需要在TensorFlow里面创建一个新的操作符(Op),并通过注册机制告诉TensorFlow,这时需要提供这个算子的名字,输入输出是什么类型的,形状长啥样等等。

最后,在Python环境下,你可以用TensorFlow提供的API来把这个自定义算子加入到你的计算图中。通过tf.load_op_library()函数加载你的库文件,然后用tf.custom_op()函数把它当作一个操作来使用。

  1. PyTorch:

对于PyTorch,其实步骤差不多,但是在PyTorch中你可以直接用Python或者C++来写你的自定义算子代码。 然后,在Python环境下,你利用torch.utils.cpp_extension来编译这段代码,生成一个扩展模块。这里同样需要告诉它算子的名字,数据类型和形状等信息。 编译好之后,就可以像导入普通模块一样,把这个扩展模块导入到PyTorch中了。

最后,使用这个自定义算子就像使用PyTorch中其他的操作一样简单。

无论是在TensorFlow还是PyTorch中添加自定义算子,过程都是:“写算子,编译,注册,使用”。

本文作者:zzw

本文链接:

版权声明:本博客所有文章除特别声明外,均采用 CC BY-SA 4.0 License 许可协议。转载请注明出处!