思源、一鸣参与

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

一行代码定义计算图,So Easy,妈妈再也不用担心我的机器学习。

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

JAXnet 是一个基于 JAX 的深度学习库,它的 API 提供了便利的模型搭建体验。相比 TensorFlow 2.0 或 PyTorch 等主流框架,JAXnet 拥有独特的优势。举个栗子,不论是 Keras 还是 PyTorch,它们建模就像搭积木一样。

然而,还有一种比搭积木更简单的方法,这就是 JAXnet 的模块化:

from jaxnet import *
net = Sequential(Dense(1024), relu, Dense(1024), relu, Dense(4), logsoftmax)

创建一个全连接网络可以直接用预定义的模块,可以说 JAXnet 定义计算图,只需一行代码就可以了。写一个神经网络,原来 So easy。

总体来说,JAXnet 主要关注的是模块化、可扩展性和易用性等几个方面:

  • 采用了不可变权重,而不是全局计算图,从而获得更强的稳健性;

  • 用于构建神经网络、训练循环、预处理、后处理等过程的 NumPy 代码经过 GPU 编译;

  • 任意模块或整个网络的正则化、重参数化都只需要一行代码;

  • 取消了全局随机状态,采用了更便捷的 Key 控制。

可扩展性

你可以使用 @parametrized 定义自己的模块,并复用其它的模块:

from jax import numpy as np
@parametrized
def loss(inputs, targets):
    return -np.mean(net(inputs) * targets)

所有的模块都是用这样的方法组合在一起的。jax.numpy (https://github.com/google/jax#whats-supported) 是 numpy 的镜像。只要你知道怎么使用 numpy,那么你就可以知道 JAXnet 大部分的用法了。

以下是 TensorFlow2/Keras 的代码,JAXnet 相比之下更为简洁:

import tensorflow as tf
from tensorflow.keras import Sequential
from tensorflow.keras.layers import Dense, Lambda

net = Sequential([Dense(1024, 'relu'), Dense(1024, 'relu'), Dense(4), Lambda(tf.nn.log_softmax)])

def loss(inputs, targets):
    return -tf.reduce_mean(net(inputs) * targets)

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

非可变权重

和 TensorFlow 或者 Keras 不同,JAXnet 没有全局计算图。net 和 loss 这样的模块不保存可变权重。权重则是保存在分开的不可变类中。这些权重由 init_parameters 函数初始化,用于提供随机的键和样本输入:

from jax.random import PRNGKey
def next_batch(): 
    return np.zeros((3, 784)), np.zeros((3, 4))
params = loss.init_parameters(PRNGKey(0), *next_batch())
print(params.sequential.dense2.bias) # [0.00376661 0.01038619 0.00920947 0.00792002]

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

opt = optimizers.Adam()
state = opt.init_state(params)for _ in range(10):
state = opt.optimize(loss.apply, state, *next_batch()) # accelerate with jit=True

trained_params = opt.get_parameters(state)

当需要对网络进行评价时:

test_loss = loss.apply(trained_params, *test_batch) # accelerate with jit=True

JAXnet 的正则化也十分简单:

loss = L2egularized(loss, scale = .1)

其他特性

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

安装也十分简单,使用 pip 安装即可。如果需要使用 GPU,则需要先安装 jaxlib。

其他具体的 API 可参考:https://github.com/JuliusKunze/jaxnet/blob/master/API.md

工程深度学习框架计算图jitJAX
1
暂无评论
暂无评论~