设为首页收藏本站

LUPA开源社区

 找回密码
 注册
文章 帖子 博客
LUPA开源社区 首页 业界资讯 开源资讯 Google 查看内容

谷歌开源TensorFlow的简化库JAX

2018-12-17 01:28| 发布者: joejoe0332| 查看: 490| 评论: 0|原作者: oschina|来自: oschina

摘要: 谷歌开源了一个 TensorFlow 的简化库 JAX。JAX 结合了 Autograd 和 XLA,专门用于高性能机器学习研究。凭借 Autograd,JAX 可以求导循环、分支、递归和闭包函数,并且它可以进行三阶求导。通过 grad,它支持自动模式 ...

谷歌开源了一个 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


酷毙

雷人

鲜花

鸡蛋

漂亮
  • 快毕业了,没工作经验,
    找份工作好难啊?
    赶紧去人才芯片公司磨练吧!!

最新评论

关于LUPA|人才芯片工程|人才招聘|LUPA认证|LUPA教育|LUPA开源社区 ( 浙B2-20090187 浙公网安备 33010602006705号   

返回顶部