内容简介:谷歌开源了一个 TensorFlow 的简化库 JAX。JAX 结合了 Autograd 和 XLA,专门用于高性能机器学习研究。
谷歌开源了一个 TensorFlow 的简化库 JAX。
JAX 结合了 Autograd 和 XLA,专门用于高性能机器学习研究。
凭借 Autograd,JAX 可以求导循环、分支、递归和闭包函数,并且它可以进行三阶求导。通过 grad,它支持自动模式反向求导(反向传播)和正向求导,且二者可以任何顺序任意组合。
得力于 XLA,可以在 GPU 和 TPU 上编译和运行 NumPy 程序。默认情况下,编译发生在底层,库调用实时编译和执行。但是 JAX 还允许使用单一函数 API jit 将 Python 函数及时编译为 XLA 优化的内核。编译和自动求导可以任意组合,因此可以在 Python 环境下实现复杂的算法并获得最大的性能。
demo:
import jax.numpy as np from jax import grad, jit, vmap from functools import partial def predict(params, inputs): for W, b in params: outputs = np.dot(inputs, W) + b inputs = np.tanh(outputs) return outputs def logprob_fun(params, inputs, targets): preds = predict(params, inputs) return np.sum((preds - targets)**2) grad_fun = jit(grad(logprob_fun)) # compiled gradient evaluation function perex_grads = jit(vmap(grad_fun, in_axes=(None, 0, 0))) # fast per-example grads
更深入地看,JAX 实际上是一个可扩展的可组合函数转换系统,grad 和 jit 都是这种转换的实例。
项目地址:https://github.com/google/JAX
以上就是本文的全部内容,希望本文的内容对大家的学习或者工作能带来一定的帮助,也希望大家多多支持 码农网
猜你喜欢:- 开源编辑器 Atom 简化代码审查过程
- Imperva 开源域目录控制器,简化活动目录集成
- 简化数据获取!Uber开源深度学习分布训练库Petastorm
- TensorFlow 官方简化版!谷歌开源机器学习库 JAX
- 灵活、轻便,阿里开源简化GNN应用框架Graph-Learn
- Jetpack Compose 开源:基于 Kotlin 的响应式编程方案,简化 UI 开发
本站部分资源来源于网络,本站转载出于传递更多信息之目的,版权归原作者或者来源机构所有,如转载稿涉及版权问题,请联系我们。
汇编语言(第3版)
王爽 / 清华大学出版社 / 2013-9 / 36.00元
《汇编语言(第3版)》具有如下特点:采用了全新的结构对课程的内容进行组织,对知识进行最小化分割,为读者构造了循序渐进的学习线索;在深入本质的层面上对汇编语言进行讲解;对关键环节进行深入的剖析。《汇编语言(第3版)》可用作大学计算机专业本科生的汇编教材及希望深入学习计算机科学的读者的自学教材。一起来看看 《汇编语言(第3版)》 这本书的介绍吧!