機(jī)器之心報(bào)道

創(chuàng)新互聯(lián)建站的客戶來(lái)自各行各業(yè)，為了共同目標(biāo)，我們?cè)诠ぷ魃厦芮信浜?，從?chuàng)業(yè)型小企業(yè)到企事業(yè)單位，感謝他們對(duì)我們的要求，感謝他們從不同領(lǐng)域給我們帶來(lái)的挑戰(zhàn)，讓我們激情的團(tuán)隊(duì)有機(jī)會(huì)用頭腦與智慧不斷的給客戶帶來(lái)驚喜。專業(yè)領(lǐng)域包括網(wǎng)站設(shè)計(jì)、網(wǎng)站制作、電商網(wǎng)站開(kāi)發(fā)、微信營(yíng)銷、系統(tǒng)平臺(tái)開(kāi)發(fā)。

參與：一鳴

Jax 是一個(gè)優(yōu)秀的代碼庫(kù)，在進(jìn)行科學(xué)計(jì)算的同時(shí)能夠自動(dòng)微分，還有 GPU、TPU 的性能加速加持。但是 Jax 的生態(tài)還不夠完善，使用者相比 TF、PyTorch 少得多。近日，DeepMind 開(kāi)源了兩個(gè)基于 Jax 的新庫(kù)，給這個(gè)生態(tài)注入了新的活力。

Jax 是谷歌開(kāi)源的一個(gè)科學(xué)計(jì)算庫(kù)，能對(duì) Python 程序與 NumPy 運(yùn)算執(zhí)行自動(dòng)微分，而且能夠在 GPU 和 TPU 上運(yùn)行，具有很高的性能。基于 Jax 已有很多優(yōu)秀的開(kāi)源項(xiàng)目，如 Trax 等。近日，DeepMind 開(kāi)源了兩個(gè)基于 Jax 的新機(jī)器學(xué)習(xí)庫(kù)，分別是 Haiku 和 RLax，它們都有著各自的特色，對(duì)于豐富深度學(xué)習(xí)社區(qū)框架、提升研究者和開(kāi)發(fā)者的使用體驗(yàn)有著不小的意義。
Haiku：https://github.com/deepmind/haikuRLax：https://github.com/deepmind/rlax
Haiku：在 Jax 上進(jìn)行面向?qū)ο箝_(kāi)發(fā)
首先值得注意的是 Haiku，這是一個(gè)面向 Jax 的深度學(xué)習(xí)代碼庫(kù)，它是由 Sonnet 作者——一個(gè)谷歌的神經(jīng)網(wǎng)絡(luò)庫(kù)團(tuán)隊(duì)開(kāi)發(fā)的。
為什么要使用 Haiku？這是因?yàn)槠渲С值氖?Jax，Jax 在靈活性和性能上具有相當(dāng)?shù)膬?yōu)勢(shì)。但是另一方面，Jax 本身是函數(shù)式的，和面向?qū)ο蟮挠脩袅?xí)慣有差別。因此，通過(guò) Haiku，用戶可以在 Jax 上進(jìn)行面向?qū)ο箝_(kāi)發(fā)了。
此外，Haiku 的 API 和編程模型都是基于 Sonnet，因此使用過(guò) Sonnet 的用戶可以快速上手。項(xiàng)目作者也表示，Sonnet 之于 TensorFlow 的提升就如同 Haiku 之于 Jax。
目前，Haiku 已公開(kāi)了 Alpha 版本，已完全開(kāi)源。項(xiàng)目作者歡迎使用者提出建議。
Haiku 怎么和 Jax 交互
Haiku 主要分為兩個(gè)模塊：hk.Modules和 hk.transform。下文將會(huì)分別介紹。
hk.Modules 是 Python 對(duì)象，保存著到參數(shù)、其他模塊和方法的參照（references）。
hk.transform 則負(fù)責(zé)將面向?qū)ο蟮哪K轉(zhuǎn)換為純粹的函數(shù)式代碼，然后讓 jax 中的 jax.jit, jax.grad, jax.pmap 等進(jìn)行處理，從而實(shí)現(xiàn)和 Jax 組件的兼容。
Haiku 的功能
Haiku 能夠做到很多機(jī)器學(xué)習(xí)需要完成的任務(wù)，相關(guān)功能和代碼如下：

自定義你的模塊

在 Haiku 中，類似于 TF2.0 和 PyTorch，你可以自定義模塊，作為 hk.Module 的子類。例如，自定義一個(gè)線性層：

class MyLinear(hk.Module):  def __init__(self, output_size, name=None):    super(MyLinear, self).__init__(name=name)    self.output_size = output_size  def __call__(self, x):    j, k = x.shape[-1], self.output_size    w_init = hk.initializers.TruncatedNormal(1. / np.sqrt(j))    w = hk.get_parameter("w", shape=[j, k], dtype=x.dtype, init=w_init)    b = hk.get_parameter("b", shape=[k], dtype=x.dtype, init=jnp.zeros)    return jnp.dot(x, w) + b

可以看出，Haiku 的代碼和 TensorFlow 等非常相似，但是你可以看到包括 numpy 等的方法還可以定義在模塊中。Haiku 的優(yōu)勢(shì)就在于，它不是一個(gè)封閉的框架，而是代碼庫(kù)，因此可以在定義模塊的過(guò)程中調(diào)用其他的庫(kù)和方法。
當(dāng)定義好線性層后，我們想要試試自動(dòng)微分的方法了：

def forward_fn(x):  model = MyLinear(10)  return model(x)# Turn `forward_fn` into an object with `init` and `apply` methods.forward = hk.transform(forward_fn)x = jnp.ones([1, 1])# When we run `forward.init`, Haiku will run `forward(x)` and collect initial# parameter values. Haiku requires you pass a RNG key to `init`, since parameters# are typically initialized randomly:key = hk.PRNGSequence(42)params = forward.init(next(key), x)# When we run `forward.apply`, Haiku will run `forward(x)` and inject parameter# values from the `params` that are passed as the first argument. We do not require# an RNG key by default since models are deterministic. You can (of course!) change# this using `hk.transform(f, apply_rng=True)` if you prefer:y = forward.apply(params, x)

這里可以看到，定義好模塊和前向傳播的函數(shù)后，使用 hk.transform(forward_fn) 可以將這種面向?qū)ο蟮姆椒ㄞD(zhuǎn)換成 Jax 底層的函數(shù)式代碼進(jìn)行處理，因此你不需要擔(dān)心底層的計(jì)算問(wèn)題。另外，這里的代碼相比 TensorFlow 還要簡(jiǎn)潔。
非訓(xùn)練狀態(tài)
有時(shí)候，我們想要在訓(xùn)練的過(guò)程中保持某些內(nèi)部參數(shù)的狀態(tài)，在 Haiku 上這也是非常容易實(shí)現(xiàn)的。

def forward(x, is_training):  net = hk.nets.ResNet50(1000)  return net(x, is_training)forward = hk.transform_with_state(forward)# The `init` function now returns parameters **and** state. State contains# anything that was created using `hk.set_state`. The structure is the same as# params (e.g. it is a per-module mapping of named values).params, state = forward.init(rng, x, is_training=True)# The apply function now takes both params **and** state. Additionally it will# return updated values for state. In the resnet example this will be the# updated values for moving averages used in the batch norm layers.logits, state = forward.apply(params, state, rng, x, is_training=True)

如上所示，只需要兩行代碼進(jìn)行設(shè)置。
和 jax.pmap 聯(lián)合進(jìn)行分布式訓(xùn)練
由于所有的代碼都會(huì)被轉(zhuǎn)換成 Jax 的函數(shù)，因此它們和 jax.pmap. 是完全兼容的。這說(shuō)明，我們可以利用 jax.pmap 來(lái)進(jìn)行分布式計(jì)算。
如下為進(jìn)行數(shù)據(jù)分割的分布式加速代碼，首先，我們先定義模型和訓(xùn)練步驟：

def loss_fn(inputs, labels):  logits = hk.nets.MLP([8, 4, 2])(x)  return jnp.mean(softmax_cross_entropy(logits, labels))loss_obj = hk.transform(loss_fn)# Initialize the model on a single device.rng = jax.random.PRNGKey(428)sample_image, sample_label = next(input_dataset)params = loss_obj.init(rng, sample_image, sample_label)

然后設(shè)定將參數(shù)拷貝到所有的設(shè)備上：

# Replicate params onto all devices.num_devices = jax.local_device_count()params = jax.tree_util.tree_map(lambda x: np.stack([x] * num_devices), params)

定義數(shù)據(jù)分批的方法，以及參數(shù)更新的方法：

def make_superbatch():  """Constructs a superbatch, i.e. one batch of data per device."""  # Get N batches, then split into list-of-images and list-of-labels.  superbatch = [next(input_dataset) for _ in range(num_devices)]  superbatch_images, superbatch_labels = zip(*superbatch)  # Stack the superbatches to be one array with a leading dimension, rather than  # a python list. This is what `jax.pmap` expects as input.  superbatch_images = np.stack(superbatch_images)  superbatch_labels = np.stack(superbatch_labels)  return superbatch_images, superbatch_labelsdef update(params, inputs, labels, axis_name='i'):  """Updates params based on performance on inputs and labels."""  grads = jax.grad(loss_obj.apply)(params, inputs, labels)  # Take the mean of the gradients across all data-parallel replicas.  grads = jax.lax.pmean(grads, axis_name)  # Update parameters using SGD or Adam or ...  new_params = my_update_rule(params, grads)  return new_params

最后開(kāi)始分布式計(jì)算即可：

# Run several training updates.for _ in range(10):  superbatch_images, superbatch_labels = make_superbatch()  params = jax.pmap(update, axis_name='i')(params, superbatch_images,                                           superbatch_labels)

RLax：Jax 上也有強(qiáng)化學(xué)習(xí)庫(kù)了
除了令人印象深刻的 Haiku 外，DeepMind 還開(kāi)源了 RLax——這是一個(gè)基于 Jax 的強(qiáng)化學(xué)習(xí)庫(kù)。
相比 Haiku，RLax 專門(mén)針對(duì)強(qiáng)化學(xué)習(xí)。項(xiàng)目作者認(rèn)為，盡管強(qiáng)化學(xué)習(xí)中的算子和函數(shù)并不是完全的算法，但是，如果需要構(gòu)建完全基于函數(shù)式的智能體，就需要特定的數(shù)學(xué)算子。
因此，函數(shù)式的 Jax 就成為了一個(gè)不錯(cuò)的選擇。在 Jax 上進(jìn)行一定的開(kāi)發(fā)后，就可以有專用的強(qiáng)化學(xué)習(xí)庫(kù)了。RLax 目前的資料還較少，但項(xiàng)目已提供了一個(gè)示例代碼：使用 RLax 進(jìn)行 Q-learning 模型的搭建和訓(xùn)練。
代碼如下，首先，使用 Haiku 構(gòu)建基本的強(qiáng)化學(xué)習(xí)模型：

def build_network(num_actions: int) -> hk.Transformed:  def q(obs):    flatten = lambda x: jnp.reshape(x, (-1,))    network = hk.Sequential(        [flatten, nets.MLP([FLAGS.hidden_units, num_actions])])    return network(obs)  return hk.transform(q)

設(shè)定訓(xùn)練的方法：

def main_loop(unused_arg):  env = catch.Catch(seed=FLAGS.seed)  rng = hk.PRNGSequence(jax.random.PRNGKey(FLAGS.seed))  # Build and initialize Q-network.  num_actions = env.action_spec().num_values  network = build_network(num_actions)  sample_input = env.observation_spec().generate_value()  net_params = network.init(next(rng), sample_input)  # Build and initialize optimizer.  optimizer = optix.adam(FLAGS.learning_rate)  opt_state = optimizer.init(net_params)

以下和 Jax 結(jié)合，定義策略、獎(jiǎng)勵(lì)等：

@jax.jitdef policy(net_params, key, obs):    """Sample action from epsilon-greedy policy."""    q = network.apply(net_params, obs)    a = rlax.epsilon_greedy(epsilon=FLAGS.epsilon).sample(key, q)    return q, a@jax.jitdef eval_policy(net_params, key, obs):    """Sample action from greedy policy."""    q = network.apply(net_params, obs)    return rlax.greedy().sample(key, q)@jax.jitdef update(net_params, opt_state, obs_tm1, a_tm1, r_t, discount_t, q_t):    """Update network weights wrt Q-learning loss."""def q_learning_loss(net_params, obs_tm1, a_tm1, r_t, discount_t, q_t):    q_tm1 = network.apply(net_params, obs_tm1)    td_error = rlax.q_learning(q_tm1, a_tm1, r_t, discount_t, q_t)    return rlax.l2_loss(td_error)    dloss_dtheta = jax.grad(q_learning_loss)(net_params, obs_tm1, a_tm1, r_t,                                             discount_t, q_t)    updates, opt_state = optimizer.update(dloss_dtheta, opt_state)    net_params = optix.apply_updates(net_params, updates)    return net_params, opt_stateprint(f"Training agent for {FLAGS.train_episodes} episodes...")

可以看到，RLax 基于 jax.jit 的方法，在性能方面有不錯(cuò)的提升。更有趣的是，構(gòu)建模型的過(guò)程中使用了前文提到的 Haiku，可見(jiàn)基于 Jax 生態(tài)的代碼庫(kù)之間都是可以兼容的。
從 DeepMind 近日開(kāi)源的兩個(gè)代碼庫(kù)可以看到，雖然現(xiàn)在深度學(xué)習(xí)框架依然在穩(wěn)步發(fā)展，但是針對(duì)高性能的科學(xué)計(jì)算也漸漸變得更為重要了。而 Jax 這樣的優(yōu)秀開(kāi)源項(xiàng)目，無(wú)疑也需要更多的生態(tài)支持。這次開(kāi)源的 Haiku 和 RLax，無(wú)疑能夠鞏固 Jax 的地位，使其優(yōu)秀的特性進(jìn)一步得到發(fā)揮。