jax JAXnet：一行代码定义计算图，兼容三大主流框架，可GPU加速

机器心脏排列

参与:思源和益铭

一行代码定义计算图，好容易，我妈再也不用担心我的机器学习了。

项目地址:https://github.com/JuliusKunze/jaxnet

JAXnet是一个基于JAX的深度学习库，它的API提供了便捷的建模体验。与主流框架如TensorFlow 2.0或PyTorch相比，JAXnet具有独特的优势。拿栗子来说，不管是Keras还是PyTorch。它们被模拟成积木。

然而，有一种比构建块更简单的方法，那就是JAXnet的模块化:

fromjaxnet import*

净=顺序，密集，密集，对数最大值)

从预定义的模块创建神经网络模型。

要创建完全连接的网络，您可以直接使用预定义的模块。可以说JAXnet定义了计算图，只需要一行代码。写一个神经网络，原来这么简单。

一般来说，JAXnet侧重于模块化、可扩展性和易用性:

采用了不可变权重，而不是全局计算图，从而获得更强的稳健性； 用于构建神经网络、训练循环、预处理、后处理等过程的 NumPy 代码经过 GPU 编译； 任意模块或整个网络的正则化、重参数化都只需要一行代码； 取消了全局随机状态，采用了更便捷的 Key 控制。

扩展性

您可以使用@ parameterized来定义自己的模块并重用其他模块:

fromjax importnumpy asnp

@ parameter ized floss:return-NP . mean* targets)

所有模块都是这样组合的。杰克斯·努皮是努皮的镜像。只要知道如何使用numpy，就可以知道JAXnet的大部分用法。

以下是TensorFlow2/Keras的代码，比JAXnet简单:

importtensorflow astffromtensorflow.keras importSequentialfromtensorflow.keras.layers importDense, Lambda

net =顺序、密集、密集、λ]

defloss:return-tf.reduce_mean*目标)

需要注意的是，Lambda函数在JAXnet中不是必须的。relu和logsoftmax函数都是Python写的。

不可变重量

不像TensorFlow或者Keras，JAXnet没有全局计算图。模块，如。净价和损失不能节省可变权重。权重存储在不同的不可变类中。这些权重由init_parameters函数初始化，以提供随机键和样本输入:

fromjax.random importPRNGKeydefnext_batch:returnnp.zeros), np.zeros)

params = loss . init _ parameters，* next _ batch)print#

目标函数不在线改变权重，而是不断更新下一个版本的权重。它们以新的优化状态返回，并由get_parameters检索。

opt = optimizers.Adamstate = opt.init_statefor _ in range:state = opt.optimize # accelerate with jit=Truetrained_params = opt.get_parameters

当网络需要评估时:

test_loss= loss.apply # accelerate with jit=True

JAXnet的正则化也很简单:

loss= L2egularized

其他特征

除了简洁的代码，JAXnet还支持GPU上的计算。也可以用jit编译，摆脱Python运行缓慢的问题。同时，JAXnet是单步调试，就像Python代码一样。

安装也很简单，用pip安装即可。如果需要使用GPU，需要先安装jaxlib。

其他具体的空气污染指数可参考:https://github.com/JuliusKunze/jaxnet/blob/master/API.md

这篇文章是为机器的核心而写的。请联系本微信官方账号进行授权。