如何在jax框架的jit中快速的实现循环结构 —— Jax框架的jit编译是否可以使用循环结构
相关:
Jax框架的jit编译是否可以使用循环结构,如果使用循环结构需要注意什么
前文中给出了jit下使用python做循环结构的代码,下面再次给出这个代码,这个代码为jupyter-notebook环境,并且在jit编译时需要60秒,运行9.4秒左右:
from jax import jit, random
import jax.numpy as jnp
from functools import partial
@partial(jit, static_argnums=(2,))
def f(x, y, z):
print("Running f():")
print(f" x = {x}")
print(f" y = {y}")
print(f" z = {z}")
for _ in range(z):
y = jnp.dot(x + 0.0001, y + 0.0001)
print(f" result = {y}")
return y
key = random.PRNGKey(0)
x = random.normal(key, (10000, 10000))
y = random.normal(key, (10000, ))
z = 10000
%timeit f(x, y, z).block_until_ready()
Jax作为TensorFlow的改进版,在一定程度上和TensorFlow的特性保持一致,比如在循环、判断这样的结构上,如果在jit中使用python做控制就会使效率比较低(当然我们可以像pytorch那边把循环结构提出来后再用python来实现,这样的效率就会几乎无差别),但是jax和TensorFlow一样提供了可以进行原生编译的循环和判断结构,下面给出jax.lax.scan来实现上面代码的循环结构的实现:
import jax.numpy as jnp
from jax import random
key = random.PRNGKey(0)
x = random.normal(key, (10000, 10000))
y = random.normal(key, (10000, ))
def body(arr, extra):
t = jnp.dot(arr[0] + 0.0001, arr[1] + 0.0001)
return (arr[0], t), extra
%timeit lax.scan(body, (x, y), jnp.ones(10000))
运行上面这个用jax.lax.scan实现的循环结构可以使用几乎不用考虑的编译时间然后直接进入运算阶段,因为这里并不需要像在jit中将python循环结构展开编译的过程,而是直接使用原生的编译好的循环操作来实现。
运行结果:
使用jax的原生循环或判断语句进行流程控制,可以避免掉大量的编译时间(比如这里的例子就避免掉了60秒的编译时间而且我们可以灵活的控制循环次数),效率更高,并且更加的灵活。但是,使用jax的原生控制语句和TensorFlow中使用TensorFlow的原生控制语句一样,都面临着增加学习成本的问题,并且到时使用难度加大,这一点估计是编译型计算框架无法摆脱的native issue问题。
比如这里给出30000次循环次数的代码:(几乎无编译时间)
如何在jax框架的jit中快速的实现循环结构 —— Jax框架的jit编译是否可以使用循环结构的更多相关文章
- 布隆过滤器 - 如何在100个亿URL中快速判断某URL是否存在?
题目描述 一个网站有 100 亿 url 存在一个黑名单中,每条 url 平均 64 字节.这个黑名单要怎么存?若此时随便输入一个 url,你如何快速判断该 url 是否在这个黑名单中? 题目解析 这 ...
- 小师妹学JVM之:JIT中的PrintCompilation
目录 简介 PrintCompilation 分析PrintCompilation的结果 总结 简介 上篇文章我们讲到了JIT中的LogCompilation,将编译的日志都收集起来,存到日志文件里面 ...
- 如何在Visual Studio 2017中使用C# 7+语法 构建NetCore应用框架之实战篇(二):BitAdminCore框架定位及架构 构建NetCore应用框架之实战篇系列 构建NetCore应用框架之实战篇(一):什么是框架,如何设计一个框架 NetCore入门篇:(十二)在IIS中部署Net Core程序
如何在Visual Studio 2017中使用C# 7+语法 前言 之前不知看过哪位前辈的博文有点印象C# 7控制台开始支持执行异步方法,然后闲来无事,搞着,搞着没搞出来,然后就写了这篇博文,不 ...
- 如何在Python中快速画图——使用Jupyter notebook的魔法函数(magic function)matplotlib inline
如何在Python中快速画图--使用Jupyter notebook的魔法函数(magic function)matplotlib inline 先展示一段相关的代码: #we test the ac ...
- 如何在ASP.NET Web站点中统一页面布局[Creating a Consistent Layout in ASP.NET Web Pages(Razor) Sites]
如何在ASP.NET Web站点中统一页面布局[Creating a Consistent Layout in ASP.NET Web Pages(Razor) Sites] 一.布局页面介绍[Abo ...
- 【转】快速理解Kafka分布式消息队列框架
from:http://blog.csdn.net/colorant/article/details/12081909 快速理解Kafka分布式消息队列框架 标签: kafkamessage que ...
- MFC中快速应用OpenCV(转)
转载链接:http://wiki.opencv.org.cn/index.php/MFC%E4%B8%AD%E5%BF%AB%E9%80%9F%E5%BA%94%E7%94%A8OpenCV 简介和缘 ...
- [.NET] 《Effective C#》快速笔记(四)- 使用框架
<Effective C#>快速笔记(四)- 使用框架 .NET 是一个类库,你了解的越多,自己需要编写的代码就越少. 目录 三十.使用重写而不是事件处理函数 三十一.使用 ICompar ...
- 在WinForm应用程序中快速实现多语言的处理(2)--开发框架模块的整合
我在上篇随笔<在WinForm应用程序中快速实现多语言的处理>里面介绍了Winform开发中多语言的处理解决方案,整个多语言解决方案是以实际需求为驱动,以减少代码改动,高效处理为目的,通过 ...
- 探讨后端选型中不同语言及对应的Web框架
在进行后端选型的时候,实际上我们要选择的是一个框架.后端领域所使用的技术和框架已经趋于稳定,我们只需要按我们的需要选择所需要的框架.当存在多个框架适合时,我们再选择适合的语言.不得不指出的是,当我们喜 ...
随机推荐
- 阻塞外挂 TCP 端口 让外挂服务器增加无用处理 反攻击 是4个IP 苹果 安卓 pc 域名
using namespace std;#include<stdlib.h>#pragma comment(lib,"WS2_32.lib") #include < ...
- 使用命令行(powershell)压缩(7Z RAR)指定日期文件
使用命令行(powershell)压缩(7Z RAR)指定日期文件,powershell ,7z. WINDOWS命令行是无法按时间过滤文件的,我们通过powershell 里的Get-ChildIt ...
- vitual box 安装centos7
vitual box下载地址: https://www.virtualbox.org/wiki/Downloads centos7下载地址: CentOS Linux 选择7-2009,x86-64 ...
- 华擎B365 BIOS 设置来电启动,来电自启,来电后开机
开机 DEL 进入BIOS. F6 进入高级模式. 在Advanced \ Chipset Configuration ,拉到底,找到 Restore on AC/Power Loss,设置为 Pow ...
- 前端使用 Konva 实现可视化设计器(14)- 折线 - 最优路径应用【代码篇】
话接上回<前端使用 Konva 实现可视化设计器(13)- 折线 - 最优路径应用[思路篇]>,这一章继续说说相关的代码如何构思的,如何一步步构建数据模型可供 AStar 算法进行路径规划 ...
- 为什么说 Mybatis 是半自动 ORM 映射工具?它与全自动的区别在哪里?
Hibernate 属于全自动 ORM 映射工具,使用 Hibernate 查询关联对象或者关联集合对象时,可以根据对象关系模型直接获取,所以它是全自动的.而 Mybatis 在查询关联对象或关联集合 ...
- Javascript高级程序设计第三章 | ch3 | 阅读笔记
语言基础 语法 标识符 注释 // /* */ 严格模式 // 也可以单独指定在一个函数中进行 'use strict' 语句 语句末尾分号不是必须的,但是最好加上 加上分号方便开发者删除空行压缩代码 ...
- 短链服务接口慢优化 redis应用
短链服务接口慢优化 redis应用 短链接服务:1.长链接 查询 短链接(长链接如果存在,直接返回短链接,如果长链接不存在,则需要生成短链接),比如:在获取短信之前,或者管理后台编辑短信内容之前,需要 ...
- Java代码忽略https证书:解决No subject alternative names present问题 HttpURLConnection https请求
Java代码忽略https证书:解决No subject alternative names present问题 import org.slf4j.Logger; import org.slf4j.L ...
- #PowerBi Superchange PowerBi 数据模型篇(1)
在过去,业务人员通常不熟悉数据建模,因为数据建模往往是专业的IT人员的领域.但是,得益于POWERBI和POWER PIVOT FOR EXCEL ,这样的日子一去不复返了.那么,什么是数据模型呢? ...