技术

上下文记忆 agentic chat 图数据库的一些考量 LLM一些探索 Agent实践 LLM预训练 向量数据库的一些考量 fastapi+sqlalchemy进行项目开发 LLM微调实践 Python协程实现 Agent Functon Calling LLamaIndex入门 Multi-Agent探索 Python虚拟机 LLM工作流编排 Python实践 下一个平台Agent 激发LLM涌现——提示工程 LLM微调理论 大佬沉思 LLM外挂知识库 LLMOps 多模态LLM Python一些比较有意思的库 Transformers源码学习 LangChain源码学习 通用分布式计算引擎Ray Python并发 go依赖注入 go collection gc的基本原理 golang性能分析及优化 数据湖 高性能计算与存储 Linux2.1.13网络源代码学习 《大数据经典论文解读》 三驾马车学习 Spark 内存管理及调优 Yarn学习 从Spark部署模式开始讲源码分析 容器狂占内存资源怎么办? 多角度理解一致性 golang io使用及优化模式 Flink学习 c++学习 学习ebpf go设计哲学 ceph学习 学习mesh kvm虚拟化 学习MQ go编译器以及defer实现 学习go 为什么要有堆栈 汇编语言 计算机组成原理 运行时和库 Prometheus client mysql 事务 mysql 事务的隔离级别 mysql 索引 坏味道 学习分布式 学习网络 学习Linux go堆内存分配 golang 系统调用与阻塞处理 Goroutine 调度过程 重新认识cpu mosn有的没的 负载均衡泛谈 单元测试的新解读 《Redis核心技术与实现》笔记 《Prometheus监控实战》笔记 Prometheus 告警学习 calico源码分析 对容器云平台的理解 Prometheus 源码分析 并发的成本 基础设施优化 hashicorp raft源码学习 docker 架构 mosn细节 与微服务框架整合 Java动态代理 编程范式 并发通信模型 《网络是怎样连接的》笔记 go channel codereview gc分析 jvm 线程实现 go打包机制 go interface及反射 如何学习Kubernetes 《编译原理之美》笔记——后端部分 《编译原理之美》笔记——前端部分 Pilot MCP协议分析 go gc 内存管理玩法汇总 软件机制 istio流量管理 Pilot源码分析 golang io 学习Spring mosn源码浅析 MOSN简介 《datacenter as a computer》笔记 学习JVM Tomcat源码分析 Linux可观测性 学习存储 学计算 Gotty源码分析 kubernetes operator kaggle泰坦尼克问题实践 kubernetes扩缩容 神经网络模型优化 直觉上理解深度学习 如何学习机器学习 TIDB源码分析 什么是云原生 Alibaba Java诊断工具Arthas TIDB存储——TIKV 《Apache Kafka源码分析》——简介 netty中的线程池 guava cache 源码分析 Springboot 启动过程分析 Spring 创建Bean的年代变迁 Linux内存管理 自定义CNI IPAM 共识算法 spring redis 源码分析 kafka实践 spring kafka 源码分析 Linux进程调度 让kafka支持优先级队列 Codis源码分析 Redis源码分析 C语言学习 《趣谈Linux操作系统》笔记 docker和k8s安全访问机制 jvm crash分析 Prometheus 学习 Kubernetes监控 Kubernetes 控制器模型 容器日志采集 容器狂占资源怎么办? Kubernetes资源调度——scheduler 时序性数据库介绍及对比 influxdb入门 maven的基本概念 《Apache Kafka源码分析》——server Kubernetes类型系统 源码分析体会 《数据结构与算法之美》——算法新解 Kubernetes源码分析——controller mananger Kubernetes源码分析——apiserver Kubernetes源码分析——kubelet Kubernetes介绍 ansible学习 Kubernetes源码分析——从kubectl开始 jib源码分析之Step实现 线程排队 jib源码分析之细节 跨主机容器通信 jib源码分析及应用 为容器选择一个合适的entrypoint kubernetes yaml配置 《持续交付36讲》笔记 mybatis学习 程序猿应该知道的 无锁数据结构和算法 CNI——容器网络是如何打通的 为什么很多业务程序猿觉得数据结构和算法没用? 串一串一致性协议 当我在说PaaS时,我在说什么 《数据结构与算法之美》——数据结构笔记 PouchContainer技术分享体会 harbor学习 用groovy 来动态化你的代码 精简代码的利器——lombok 学习 《深入剖析kubernetes》笔记 编程语言那些事儿 rxjava3——背压 rxjava2——线程切换 spring cloud 初识 《深入拆解java 虚拟机》笔记 《how tomcat works》笔记 hystrix 学习 rxjava1——概念 Redis 学习 TIDB 学习 如何分发计算 Storm 学习 AQS1——论文学习 Unsafe Spark Stream 学习 linux vfs轮廓 《自己动手写docker》笔记 java8 实践 中本聪比特币白皮书 细读 区块链泛谈 比特币 大杂烩 总纲——如何学习分布式系统 hbase 泛谈 forkjoin 泛谈 看不见摸不着的cdn是啥 《jdk8 in action》笔记 程序猿视角看网络 bgp初识 calico学习 AQS——粗略的代码分析 我们能用反射做什么 web 跨域问题 《clean code》笔记 《Elasticsearch权威指南》笔记 mockito简介及源码分析 2017软件开发小结—— 从做功能到做系统 《Apache Kafka源码分析》——clients dns隐藏的一个坑 《mysql技术内幕》笔记 log4j学习 为什么netty比较难懂? 递归、回溯、动态规划 apollo client源码分析及看待面向对象设计 学习并发 docker运行java项目的常见问题 OpenTSDB 入门 spring事务小结 分布式事务 javascript应用在哪里 《netty in action》读书笔记 netty对http2协议的解析 ssl证书是什么东西 http那些事 苹果APNs推送框架pushy apple 推送那些事儿 编写java框架的几大利器 java内存模型和jvm内存布局 java exception Linux IO学习 netty内存管理 测试环境docker化实践 netty在框架中的使用套路 Nginx简单使用 《Linux内核设计的艺术》小结 Go并发机制及语言层工具 Linux网络源代码学习——数据包的发送与接收 《docker源码分析》小结 docker namespace和cgroup zookeeper三重奏 数据库的一些知识 Spark 泛谈 链式处理的那些套路 netty回顾 Thrift基本原理与实践(二) Thrift基本原理与实践(一) 回调 异步执行抽象——Executor与Future Docker0.1.0源码分析 java gc Jedis源码分析 深度学习泛谈 Linux网络命令操作 JTA与TCC 换个角度看待设计模式 Scala初识 向Hadoop学习NIO的使用 以新的角度看数据结构 并发控制相关的硬件与内核支持 systemd 简介 quartz 源码分析 基于docker搭建测试环境(二) spring aop 实现原理简述 自己动手写spring(八) 支持AOP 自己动手写spring(七) 类结构设计调整 分析log日志 自己动手写spring(六) 支持FactoryBean 自己动手写spring(九) 总结 自己动手写spring(五) bean的生命周期管理 自己动手写spring(四) 整合xml与注解方式 自己动手写spring(三) 支持注解方式 自己动手写spring(二) 创建一个bean工厂 自己动手写spring(一) 使用digester varnish 简单使用 关于docker image的那点事儿 基于docker搭建测试环境 分布式配置系统 JVM执行 git maven/ant/gradle/make使用 再看tcp kv系统 java nio的多线程扩展 《Concurrency Models》笔记 回头看Spring IOC IntelliJ IDEA使用 Java泛型 vagrant 使用 Go常用的一些库 Python初学 Goroutine 调度模型 虚拟网络 《程序员的自我修养》小结 Kubernetes存储 访问Kubernetes上的Service Kubernetes副本管理 Kubernetes pod 组件 Go基础 JVM类加载 硬币和扑克牌问题 LRU实现 virtualbox 使用 ThreadLocal小结 docker快速入门

架构

bert rerank微调 大模型推理tips RAG向量检索与微调 dddfirework源码分析 RAG与知识图谱 大模型推理服务框架vLLM 大模型推理服务框架 模型服务化(未完成) 大模型Post-Training 大模型训练 大模型推理 从Attention到Transformer k8s设备管理 ddd从理念到代码 如何应用LLM 小鼠如何驾驭大象(LLM)? 多类型负载协调员Koordinator controller-runtime细节分析 finops学习 kubevela多集群 kubevela中cue的应用 基于k8s的工作流 kubevela源码分析 容器和CPU那些事儿 数据集管理fluid 应用管理平台kubevela karmada支持crd 多集群管理 AutoML和AutoDL 特征平台 实时训练 分布式链路追踪 K8S YAML 资源清单管理方案 tensorflow原理——python层分析 如何学习tensorflow 数据并行——allreduce 数据并行——ps 推荐系统embedding原理及实践 机器学习中的python调用c 机器学习训练框架概述 tensornet源码分析 大模型训练和推理 X的生成——特征工程 tvm tensorflow原理——core层分析 模型演变 《深度学习推荐系统实战》笔记 keras 和 Estimator tensorflow分布式训练 分布式训练的一些问题 基于Volcano的弹性训练 图神经网络 pytorch弹性分布式训练 从混部到统一调度 从RNN到Attention pytorch分布式训练 CNN 《动手学深度学习》笔记 pytorch与线性回归 多活 volcano特性源码分析 推理服务 kubebuilder 学习 mpi 学习pytorch client-go学习 提高gpu 利用率 GPU与容器的结合 GPU入门 AI云平台梳理 tensorflow学习 tf-operator源码分析 k8s批处理调度/Job调度 喜马拉雅容器化实践 Kubernetes 实践 学习rpc BFF openkruise学习 可观察性和监控系统 基于Kubernetes选主及应用 《许式伟的架构课》笔记 Admission Controller 与 Admission Webhook 发布平台系统设计 k8s水平扩缩容 Scheduler如何给Node打分 Scheduler扩展 深入controller openkruise cloneset学习 controller-runtime源码分析 pv与pvc实现 csi学习 client-go informer源码分析 kubelet 组件分析 调度实践 Pod是如何被创建出来的? 《软件设计之美》笔记 mecha 架构学习 Kubernetes events学习及应用 CRI——kubelet与容器引擎之间的接口 资源调度泛谈 业务系统设计原则 grpc学习 元编程 以应用为中心 istio学习 下一代微服务Service Mesh 《实现领域驱动设计》笔记 概率论 serverless 泛谈 《架构整洁之道》笔记 处理复杂性 那些年追过的并发 服务器端编程 网络通信协议 架构大杂烩 如何学习架构 《反应式设计模式》笔记 项目的演化特点 反应式架构摸索 函数式编程的设计模式 服务化 ddd反模式——CRUD的败笔 研发效能平台 重新看面向对象设计 业务系统设计的一些体会 函数式编程 《左耳听风》笔记 业务程序猿眼中的微服务管理 DDD实践——CQRS 项目隔离——案例研究 《编程的本质》笔记 系统故障排查汇总及教训 平台支持类系统的几个点 代码腾挪的艺术 abtest 系统设计汇总 《从0开始学架构》笔记 初级权限系统设计 领域驱动理念 现有上传协议分析 移动网络下的文件上传要注意的几个问题 推送系统的几个基本问题 做配置中心要想好的几个基本问题 不同层面的异步 分层那些事儿 性能问题分析 用户认证问题 资源的分配与回收——池 消息/任务队列

标签

k8s设备管理 多类型负载协调员Koordinator controller-runtime细节分析 finops学习 kubevela多集群 kubevela中cue的应用 基于k8s的工作流 kubevela源码分析 容器和CPU那些事儿 数据集管理fluid 应用管理平台kubevela karmada支持crd 多集群管理 K8S YAML 资源清单管理方案 从混部到统一调度 volcano特性源码分析 kubebuilder 学习 client-go学习 tf-operator源码分析 k8s批处理调度/Job调度 喜马拉雅容器化实践 Kubernetes 实践 openkruise学习 基于Kubernetes选主及应用 Admission Controller 与 Admission Webhook k8s水平扩缩容 Scheduler如何给Node打分 Scheduler扩展 深入controller openkruise cloneset学习 controller-runtime源码分析 pv与pvc实现 csi学习 client-go informer源码分析 kubelet 组件分析 调度实践 Pod是如何被创建出来的? Kubernetes events学习及应用 CRI——kubelet与容器引擎之间的接口 资源调度泛谈 如何学习Kubernetes 以应用为中心 kubernetes operator kubernetes扩缩容 serverless 泛谈 什么是云原生 自定义CNI IPAM docker和k8s安全访问机制 Kubernetes监控 Kubernetes 控制器模型 Kubernetes资源调度——scheduler Kubernetes类型系统 Kubernetes源码分析——controller mananger Kubernetes源码分析——apiserver Kubernetes源码分析——kubelet Kubernetes介绍 Kubernetes源码分析——从kubectl开始 kubernetes yaml配置 CNI——容器网络是如何打通的 当我在说PaaS时,我在说什么 《深入剖析kubernetes》笔记 Kubernetes存储 访问Kubernetes上的Service Kubernetes副本管理 Kubernetes pod 组件
上下文记忆 agentic chat bert rerank微调 大模型推理tips LLM一些探索 Agent实践 LLM预训练 RAG向量检索与微调 LLM微调实践 RAG与知识图谱 大模型推理服务框架vLLM Agent Functon Calling LLamaIndex入门 Multi-Agent探索 LLM工作流编排 大模型推理服务框架 模型服务化(未完成) 大模型Post-Training 大模型训练 大模型推理 从Attention到Transformer 下一个平台Agent 激发LLM涌现——提示工程 LLM微调理论 大佬沉思 LLM外挂知识库 LLMOps 多模态LLM Transformers源码学习 LangChain源码学习 如何应用LLM 小鼠如何驾驭大象(LLM)? AutoML和AutoDL 特征平台 实时训练 tensorflow原理——python层分析 如何学习tensorflow 数据并行——allreduce 数据并行——ps 推荐系统embedding原理及实践 机器学习中的python调用c 机器学习训练框架概述 tensornet源码分析 大模型训练和推理 X的生成——特征工程 tvm tensorflow原理——core层分析 模型演变 《深度学习推荐系统实战》笔记 keras 和 Estimator tensorflow分布式训练 分布式训练的一些问题 基于Volcano的弹性训练 图神经网络 pytorch弹性分布式训练 从RNN到Attention pytorch分布式训练 CNN 《动手学深度学习》笔记 pytorch与线性回归 推理服务 mpi 学习pytorch 提高gpu 利用率 GPU与容器的结合 GPU入门 AI云平台梳理 tensorflow学习 kaggle泰坦尼克问题实践 神经网络模型优化 概率论 直觉上理解深度学习 如何学习机器学习 深度学习泛谈

机器学习中的python调用c

2022年03月02日

简介

作为一种解释型的语言,Python的速度并不算慢。如果对速度有很高的要求的话,可以选择用更快的语言实现,比如C或C++,然后用Python调用。Python的一种常见应用场景是实现高级的逻辑。Python的解释器就是用C语言写的,即CPython。解释器将Python转换成一种中间语言,叫做Python字节码,类似于汇编语言,但是包含一些更高级的指令。当一个运行一个Python程序的时候,评估循环不断将Python字节码转换成机器码。解释型语言的好处是方便编程和调试,但是程序的运行速度慢。其中的一种解决办法是,用C语言实现一些第三方的库,然后在Python中使用。另一种方法是使用即时编译器来替换Cpython,例如PyPy,PyPy对代码生成和Python的运行速度做了优化。

python 是解释型语言,性能是瓶颈。实现混合编程的方式

  1. 使用ctypes 库加载c++编写的动态链接库
  2. 使用pybind 将c++编译为python库
  3. 使用pythran库 将python直接转换为C++代码

跨语言调用,有几个问题

  1. 一些要素 在跨语言间的对应关系,比如c++ 的对象、函数、全局变量。c++ 函数 暴露到 python 成为 一个模块下的函数,c++ 对象则暴露为 一个python 对象。 比如在jna 中,c++ 函数会被暴露为 一个接口的方法,jna 负责提供这个接口的实现。
  2. 常见的类型(作为参数的时候) 如何跨语言打通,比如string、vector、map以及指针

tf 从swig 切到了pybind11。

双引擎 GPU 容器虚拟化,用户态和内核态的技术解析和实践分享典型的 AI 软硬件生态都分为这样几个层次 ——应用 & 框架层,运行时层,驱动层,硬件层。最上层是用户的应用,这里包含了各种常见的框架 PaddlePaddle、TensorFlow、PyTorch 等等。在应用层之下是硬件提供商封装的 API 接口层:包含各类常用算子库与硬件运行时访问接口。

swig

tf 早期通过swig 实现python 调用c

  1. 在 pywrap_tensorflow_internal.cc 的实现中,静 态注册了一个函数符号表,实现了 Python 函数名到 C 函数名的二元关系。
  2. _pywrap_tensorflow_internal.so 包 含了整个 TensorFlow 运行时的所有符号。
  3. pywrap_tensorflow_internal.py 模块首次被导入时,自动地加载 _pywrap_tensorflow_internal.so 的动态链接库
  4. 在运行时,按 照 Python 的函数名称,匹配找到对应的 C 函数实现,最终实现 Python 到 c_api.c 具体 实现的调用关系。c_api.h 是 TensorFlow 的后端执行系统面向前端开放的公共 API 接口。

Client 存在部分 C++ 实现,即 tensorflow::Session。其中,tf.Session 实例直接持有 tensorflow::Session 实例的句柄。一般地,用户使用的是 tf.Session 实施编程

pybind11

动手学深度学习框架(2)- python 端如何调用 c++ 的代码

pybind11 是一个轻量级的只包含头文件(header-only)的 c++ 库,用于将 c++ 代码暴露给 python 调用(反之亦可,但主要还是前者)。

#include <pybind11/pybind11.h>          // pybind11的头文件
// PYBIND11_MODULE 是一个宏,实现一个 Python 扩展模块
PYBIND11_MODULE(pydemo, m){             // 定义Python模块pydemo,之后在 Python 脚本里必须用这个名字才能 import。m其实是 pybind11::module 的一个实例对象,它只是个普通的变量,起什么名字都可以,但为了写起来方便,一般都用“m”。
    m.doc() = "pybind11 demo doc";      // 模块的说明文档
    m.def("add", [](int a, int b) -> int { return a + b; });    // def函数,传递一个 Python 函数名和 C++ 的函数、函数对象或者是 lambda 表达式
}                                       // Python模块定义结束

假设这个 C++ 源文件名是“pybind.cpp”,用 g++ 把它编译成在 Python 里调用的模块,生成一个大概这样的文件:pydemo.cpython-35m-x86_64-linux-gnu.so

g++ pybind.cpp               \                  #编译的源文件
   -std=c++11 -shared -fPIC   \                 #编译成动态库
  `python3 -m pybind11 --includes` \            #获得 pybind11 所在的包含路径,让 g++ 能够找得到头文件
  -o pydemo`python3-config --extension-suffix`  #生成的动态库名字,前面必须是源码里的模块名,而后面那部分则是 Python 要求的后缀名

之后就可以在python 中使用

import pydemo
x = pydemo.add(1,2)
print(x)

进阶:C++ 里的类也能够等价地转换到 Python 里面调用

算子设计

PS:其实如果使用c++ 来写推理或训练引擎的话,就没有python调用c这个复杂的事儿了。对于一个推理框架,大概可以理解为,

  1. 先基于onnx/pnnx等模型文件,自己提一套抽象比如RuntimeGraph等将模型权重、参数加载进来,然后按拓扑排序执行,执行到某个节点时,调用其对应的算子(为此有一个全局的算子注册机制),节点(Node或Operator)为算子准备入参、拿到出参。概念上从大到下是Graph ==> node/op ==> cuda 函数。
  2. 专用的推理框架入口是onnx/pnnx等模型文件,只需要graph、节点/等概念,不需要pytorch 中类似layer概念(那是为了编程上抽象复用的)。
  3. tensor/显存的申请、释放都是上层组件负责,会有一个DeviceAllocator(分别对应cpu和gpu)组件负责内存和显存的分配和释放、内存和显存之间的copy等接口(比如tensor.to_cuda。再复杂一点先提前申请一个大的,内部再复用一下),对DeviceAllocator封装后提供tensor对象(tensor持有DeviceAllocator 引用,初始化时调用DeviceAllocator.allocate,析构时调用DeviceAllocator.release)。只是给算子函数传入input/weight/output 指针,算子也分为cpu和gpu实现。 在深度学习中,算子通常指的是在神经网络中对数据执行数学运算的函数。这些运算可以是简单的,如加法、乘法,也可以是复杂的,如卷积、池化、归一化等。根据算子内部参数的有无,我们大致可以将算子分为两大类:
  4. 带参数的,例如卷积算子,全连接算子,rmsnorm算子等。PS: 加载模型的一个重要的活儿就是用模型权重去初始化各类带参数算子。
  5. 不带参数的,例如sigmoid算子,softmax算子等

算子基类

class BaseOP {
 public:
  explicit BaseOP(base::DeviceType device_type, OPType op_type,base::DataType data_type, std::string op_name = "");
  base::DataType data_type() const;
  OPType op_type() const;
  ...
  ...
  const std::string& get_op_name() const; // 返回算子的名字
  void set_op_name(const std::string& op_name); // 设置算子的名称
  base::DeviceType device_type() const; // 返回层的设备类型
  void set_device_type(base::DeviceType device_type);
  virtual base::Status forward() = 0;

 protected:
  std::string op_name_; // 算子名
  OPType OP_type_ = OPType::kOPUnknown;             // 算子类型
  base::DataType data_type_ = base::DataType::kDataTypeUnknown; // 数据类型 fp32 或fp16 或int8
  base::DeviceType device_type_ = base::DeviceType::kDeviceUnknown;  // 设备类型 cpu或gpu
};

不带参(权重)算子类的设计

class OP : public BaseOP {
 public:
  explicit OP(base::DeviceType device_type, OPType op_type,std::string op_name = "");

  void set_input(int32_t idx, const tensor::Tensor& input) override; // 传入输入 ,需要指定这是该算子的第几个(idx)输入
  void set_output(int32_t idx, const tensor::Tensor& output) override; // 传入输出
  const tensor::Tensor& get_input(int32_t idx) const override; // 获取输入
  const tensor::Tensor& get_output(int32_t idx) const override; // 获取输出
  // 关于算子输入、输出张量的辅助函数
  size_t input_size() const override; // 获取输入的个数
  size_t output_size() const override; // 获取输出的个数
  void reset_input_size(size_t size);
  void reset_output_size(size_t size);

  virtual void to_cuda();

 private:
  std::vector<tensor::Tensor> inputs_;  // 存放输入的数组
  std::vector<tensor::Tensor> outputs_; // 存放输出的数组
};

base::Status VecAddOP::forward() {
  auto status = this->check();
  if (!status) {
    return status;
  }
  auto input1 = this->get_input(0);
  auto input2 = this->get_input(1);
  auto output = this->get_output(0);
  if (device_type_ == base::DeviceType::kDeviceCUDA) {
    CHECK(cuda_config_ != nullptr);
  }
  kernel::get_add_kernel(device_type_)(input1, input2, output,cuda_config_ ? cuda_config_->stream : nullptr);                   
  return base::error::Success();
}

带参数的算子类,多了一个类内变量用于存储权重张量

class OPFp32Param : public OP {
 public:
  explicit OPFp32Param(base::DeviceType device_type, OPType op_type,std::string op_name = "");
                    
  size_t weight_size() const;
  void reset_weight_size(size_t size);
  tensor::Tensor& get_weight(int32_t idx);
  const tensor::Tensor& get_weight(int32_t idx) const;

  void set_weight(int32_t idx, const tensor::Tensor& weight); // load model时设置权重便是靠set_weight
  void set_weight(int32_t idx, const std::vector<int32_t>& dims, const float* weight_ptr,base::DeviceType device_type = base::DeviceType::kDeviceUnknown);
       
 private:
  std::vector<tensor::Tensor> weights_; // 用于额外存放权重数据
};

base::Status RmsNormOP::forward() { // 计算的时候
  auto status = check();
  if (!status) {
    return status;
  }
  auto input = this->get_input(0);
  auto weight = this->get_weight(0);
  auto output = this->get_output(0);
  // 得到一个具体的算子计算实现
  kernel::get_rmsnorm_kernel(device_type_)(input, weight, output,cuda_config_ ? cuda_config_->stream : nullptr);                           
  return base::error::Success();
}

所谓的load_model,一般先读取模型配置(或者是一个config.json 或者是model.bin 的前xx个字节,看模型格式?),这样知道模型有几层,hidden_dim 是多少,然后才是读取每层的权重(比如每层的weight 是顺序排在model.bin中,此时顺序读即可),最终将weight数据赋值给OP.weight(一个tensor对象)。

TensorFlow自定义算子

TensorFlow 模型准实时更新上线的设计与实现计算图结构由模型的算法结构决定,对数据的操作即为 operation( op )。当模型结构确定的情况下,我们的增强就需要对 op 进行定制。 PS:介绍了针对 embedding 参数的特点,如何通过自定义op 对其进行优化。

tensorflow:自定义op简单介绍

一个Op可以接收一个或者多个输入Tensor,然后产生零个或者多个输出Tensor,分别利用Input和Output定义。在注册一个Op之后,就需要继承OpKernel,实现他的计算过程Compute函数,在Compute函数中,我们可以通过访问OpKernelContext来获得输入和输出信息。当我们需要申请新的内存地址时,可以通过OpKernelContext去申请TempTensor或者PersistentTensor。一般Tensorflow的Op都采用Eigen来操作Tensor

Adding a New Op对于 TensorFlow,可以自定义 Operation,即如果现有的库没有涵盖你想要的操作, 你可以自己定制一个。为了使定制的 Op 能够兼容原有的库,你必须做以下工作:

  1. 在一个 C++ 文件中注册新 Op. Op 的注册与实现是相互独立的. 在其注册时描述了 Op 该如何执行. 例如, 注册 Op 时定义了 Op 的名字, 并指定了它的输入和输出。
      // 最终Op被注册到了一个static变量global_op_registry中
      REGISTER_OP("ZeroOut")
     .Input("to_zero: int32")
     .Output("zeroed: int32")
     .SetShapeFn([](::tensorflow::shape_inference::InferenceContext* c) {
       c->set_output(0, c->input(0));
       return Status::OK();
     });
    
  2. 使用 C++ 实现 Op. 每一个实现称之为一个 “kernel”, 可以存在多个 kernel, 以适配不同的架构 (CPU, GPU 等)或不同的输入/输出类型.
  3. bazel(tf编译工具) 会检索所有op 并创建一个 Python wrapper. 这个wrapper是创建 Op 的公开 API. 当注册 Op 时, 会自动生成一个默认 默认的包装器. 既可以直接使用默认包装器, 也可以添加一个新的包装器.
  4. (可选) 写一个函数计算 Op 的梯度,在Python 中注册.
      @ops.RegisterGradient("ZeroOut")
      def _zero_out_grad(op, grad):
     xxxxxxxxx
    
  5. (可选) 写一个函数, 描述 Op 的输入和输出 shape. 该函数能够允许从 Op 推断 shape.
  6. 测试 Op, 通常使用 Pyhton。如果你定义了梯度,你可以使用Python的GradientChecker来测试它。

There are two main mechanisms for op and kernel registration:

  1. Static linking into the core TensorFlow library, and static initialization.
  2. Dynamic linking at runtime, using the tf.load_op_library() function. 读取op 对应的 python wrapper文件 作为python module 注册到python module中

tensorflow/custom-op

Bazel BUILD文件如下,执行 bazel build ${BAZEL_ARGS[@]} 可以得到 tensorflow/core/user_ops/zero:zero_out.so PS: 类似于执行了 上文中的g++ 编译得到so 文件。

load("//tensorflow:tensorflow.bzl", "tf_custom_op_library")
tf_custom_op_library(
    name = "zero_out.so",       #  target name
    srcs = ["zero_out.cc"],     #  the list of the sources to compile,
)

得到so 文件后,tf.load_op_library 动态加载so作为 module 使用(可以参考python module 动态加载加载)。

import tensorflow as tf
# 返回一个 A python module, containing the (op对应的)Python wrappers for Ops defined in the plugin.
# Python Module,是一个 Python 文件,以 .py 结尾,包含了 Python 对象定义和Python语句
zero_out_module = tf.load_op_library('zero_out.so')
with tf.Session():
  print(zero_out_module.zero_out([1,2,3,4,5])).eval() # eval 底层执行 session.run

zero_out_module.zero_out 可能仅用于演示,正规的ops 比如tf.matmul 等实现(对应到 编译tf生成的代码 gen_math_ops.py)会涉及到生成opDef (graphDef 的一部分)等逻辑。

tensornet框架 自定义ops 示例

  1. 运行时访问链路: import gen_xx_ops.py ==> gen_xx_ops.opxx ==> _op_def_library._apply_op_helper ==> 向graph中添加对应名字的Op节点
  2. gen_xx_ops 生成:bazel build ==> 使用 python_op_gen_main 生成 gen_xx_ops.py。PS: 有点系统调用的意思,用户态存在一个对系统调用的 封装(glibc 函数),比如调用read 其实只是传递了一个read 系统调用号,要靠内核去执行真正的read函数。
tensornet
    /core
        /kernels
            /sparse_table_ops.cc   # kernel实现
        /ops
            /sparse_table_ops.cc   # REGISTER_OP
        /BUILD                      
    /tensornet
        /core
            /gen_sparse_table_ops.py    # Bazel tf_gen_op_wrapper_py生成

Python的调用方式

# tensornet/layers/embedding_features.py
from tensornet.core import gen_sparse_table_ops
pulled_mapping_values = gen_sparse_table_ops.sparse_table_pull(...)

gen_sparse_table_ops 中的定义

# tensornet/core/gen_sparse_table_ops.py
def sparse_table_pull(resources, values, table_handle, name=None):
    ...
    _, _, _op, _outputs = _op_def_library._apply_op_helper("SparseTablePull", resources=resources, values=values,
        table_handle=table_handle, name=name)
    ...
    return _op
    ...

_op_def_library._apply_op_helper的作用是在graph中添加对应名字的Op节点。需要注意的是,Op的梯度计算节点并不是在这里加入到graph中的,这里仅仅加入了前向计算节点。

gen_sparse_table_ops.py 文件是在bazel构建过程中生成的,BUILD 文件内容

// 生成  生成op的lib,即so文件
cc_library(
    name = "sparse_table_ops_kernels",
    srcs = [
        "kernels/sparse_table_ops_dummy.cc",
        "ops/sparse_table_ops.cc",
    ],
    hdrs = [
        "//core/utility:semaphore",
    ],
    linkstatic = 1,
    deps = [
        "@org_tensorflow//tensorflow/core:framework",
        "@org_tensorflow//tensorflow/core:lib",
        "@org_tensorflow//tensorflow/core:protos_all_cc",
    ],
    alwayslink = 1,
)
// 生成python的接口,即gen_sparse_table_ops.py 文件
tf_gen_op_wrapper_py(
    name = "sparse_table_ops",
    deps = [":sparse_table_ops_kernels"],
    cc_linkopts = ['-lrt']
)

Op定义分析python_op_gen_main(tensorflow/python/framework/python_op_gen_main.cc)工具通过链接对应的so,得到对应的OpRegistry,从而生成对应的gen_xx_ops.py文件。

tensorflow c++ op 生成 python调用接口

TensorFlow 模型准实时更新上线的设计与实现定制好 op 后,如何替换模型计算图中原生的 op 呢?TensorFlow 在模型保存时,会生成 meta_graph_def 文件,文件内容是采用类似 json 的格式描述计算图的结构关系。当加载此文件时,TensorFlow 会根据文件中描述的结构信息构建出计算图。可以修改模型保存的 meta_graph_def 文件,将其中的 op 替换为我们定制的 op,同时修改每个 node 的 input 和 output 关系,以修改 op 之间的依赖关系。PS: 当然这里说的替换原有的op

horovod

很多机器学习框架都会采用如下套路:shell脚本(可选),python端 和 C++端。

  1. Shell脚本是启动运行的入口,负责解析参数,确认并且调用训练程序;
  2. Python是用户的接口,引入了C++库,封装了API,负责运行时和底层C++交互;
  3. C++实现底层训练逻辑;

深度学习分布式训练框架 horovod (2) — 从使用者角度切入

引入库的作用是获取到 C++ 的函数,并且用 python 封装一下,这样就可以在 python 世界使用 C++代码了。比如下文,python 的 _allreduce 函数就会把功能转发给 C++,由 MPI_LIB.horovod_allreduce 完成。

def _allreduce(tensor, name=None, op=Sum, prescale_factor=1.0, postscale_factor=1.0,
               ignore_name_scope=False):
    if name is None and not _executing_eagerly():
        name = 'HorovodAllreduce_%s' % _normalize_name(tensor.name)
    return MPI_LIB.horovod_allreduce(tensor, name=name, reduce_op=op,
                                     prescale_factor=prescale_factor,
                                     postscale_factor=postscale_factor,
                                     ignore_name_scope=ignore_name_scope)
## 初始化时执行
def _load_library(name):
    """Loads a .so file containing the specified operators.
    """
    filename = resource_loader.get_path_to_datafile(name)
    library = load_library.load_op_library(filename)
    return library

# Check possible symbol not found error from tensorflow version mismatch
try:
    MPI_LIB = _load_library('mpi_lib' + get_ext_suffix())
except Exception as e:
    check_installed_version('tensorflow', tf.__version__, e)
    raise e
else:
    check_installed_version('tensorflow', tf.__version__)