伯克利AI分散式框架Ray,相容TensorFlow、PyTorch與MXNet

李澤南發表於2018-01-10

不久之前,機器之心推薦了一篇論文,介紹 UC Berkeley 研究員釋出的分散式系統 Ray。開發者稱,Ray 專門為人工智慧應用設計,通過這款框架,執行於膝上型電腦上的原型演算法僅需加入數行程式碼就可以轉化為高效的分散式計算應用。近日,該框架已被開源。在本文中,伯克利官方 AI 部落格對開源框架 Ray 做了詳細介紹。

GitHub 連結:https://github.com/ray-project/ray

隨著機器學習演算法和技術的進步,出現了越來越多需要在多臺機器平行計算的機器學習應用。然而,在叢集計算裝置上執行的機器學習演算法目前仍是專門設計的。儘管對於特定的用例而言(如引數伺服器或超引數搜尋),這些解決方案的效果很好,同時 AI 領域之外也存在一些高質量的分散式系統(如 Hadoop 和 Spark),但前沿開發者們仍然常常需要從頭構建自己的系統,這意味著需要耗費大量時間和精力。

例如,應用一個簡單概念的演算法,如在強化學習中的進化策略(論文《Evolution Strategies as a Scalable Alternative to Reinforcement Learning》)。演算法包含數十行虛擬碼,其中的 Python 實現也並不多。然而,在較大的機器或叢集上執行它需要更多的軟體工程工作。作者的實現包含了上千行程式碼,以及必須定義的通訊協議、資訊序列化、反序列化策略,以及各種資料處理策略。

Ray 的目標之一在於:讓開發者可以用一個執行在膝上型電腦上的原型演算法,僅需新增數行程式碼就能輕鬆轉為適合於計算機叢集執行的(或單個多核心計算機的)高效能分散式應用。這樣的框架需要包含手動優化系統的效能優勢,同時又不需要使用者關心那些排程、資料傳輸和硬體錯誤等問題。

開源的人工智慧框架

與深度學習框架的關係:Ray 與 TensorFlow、PyTorch 和 MXNet 等深度學習框架互相相容,在很多應用上,在 Ray 中使用一個或多個深度學習框架都是非常自然的(例如,UC Berkeley 的強化學習庫就用到了很多 TensorFlow 與 PyTorch)。

與其他分散式系統的關係:目前的很多流行分散式系統都不是以構建 AI 應用為目標設計的,缺乏人工智慧應用的相應支援與 API,UC Berkeley 的研究人員認為,目前的分散式系統缺乏以下一些特性:

  • 支援毫秒級的任務處理,每秒處理百萬級的任務;
  • 巢狀並行(任務內並行化任務,例如超引數搜尋內部的並行模擬,見下圖);
  • 在執行時動態監測任意任務的依賴性(例如,忽略等待慢速的工作器);
  • 在共享可變的狀態下執行任務(例如,神經網路權重或模擬器);
  • 支援異構計算(CPU、GPU 等等)。

伯克利AI分散式框架Ray,相容TensorFlow、PyTorch與MXNet


一個巢狀並行的簡單例子。一個應用執行兩個並行實驗(每個都是長時間執行任務),每個實驗執行一定數量的並行模擬(每一個同時也是一個任務)。

Ray 有兩種主要使用方法:通過低階 API 或高階庫。高階庫是構建在低階 API 之上的。目前它們包括 Ray RLlib,一個可擴充套件強化學習庫;和 Ray.tune,一個高效分散式超引數搜尋庫。

Ray 的低層 API

開發 Ray API 的目的是讓我們能更自然地表達非常普遍的計算模式和應用,而不被限制為固定的模式,就像 MapReduce 那樣。

動態任務圖

Ray 應用的基礎是動態任務圖。這和 TensorFlow 中的計算圖很不一樣。TensorFlow 的計算圖用於表徵神經網路,在單個應用中執行很多次,而 Ray 的任務圖用於表徵整個應用,並僅執行一次。任務圖對於前臺是未知的,隨著應用的執行而動態地構建,且一個任務的執行可能建立更多的任務。

伯克利AI分散式框架Ray,相容TensorFlow、PyTorch與MXNet


計算圖示例。白色橢圓表示任務,藍色方框表示物件。箭頭表示任務依賴於一個物件,或者任務建立了一個物件。

任意的 Python 函式都可以當成任務來執行,並且可以任意地依賴於其它任務的輸出。如下示例程式碼所示:

# Define two remote functions. Invocations of these functions create tasks
# that are executed remotely.
@ray.remote
def multiply(x, y):
 return np.dot(x, y)
@ray.remote
def zeros(size):
 return np.zeros(size)
# Start two tasks in parallel. These immediately return futures and the
# tasks are executed in the background.
x_id = zeros.remote((100, 100))
y_id = zeros.remote((100, 100))
# Start a third task. This will not be scheduled until the first two
# tasks have completed.
z_id = multiply.remote(x_id, y_id)
# Get the result. This will block until the third task completes.
z = ray.get(z_id)

動作器(Actor)

僅用遠端函式和上述的任務所無法完成的一件事是在相同的共享可變狀態上執行多個任務。這在很多機器學習場景中都出現過,其中共享狀態可能是模擬器的狀態、神經網路的權重或其它。Ray 使用 actor 抽象以封裝多個任務之間共享的可變狀態。以下是關於 Atari 模擬器的虛構示例:

import gym
@ray.remote
class Simulator(object):
 def __init__(self):
     self.env = gym.make("Pong-v0")
     self.env.reset()
 def step(self, action):
     return self.env.step(action)
# Create a simulator, this will start a remote process that will run
# all methods for this actor.
simulator = Simulator.remote()
observations = []
for _ in range(4):
 # Take action 0 in the simulator. This call does not block and
 # it returns a future.
 observations.append(simulator.step.remote(0))

Actor 可以很靈活地應用。例如,actor 可以封裝模擬器或神經網路策略,並且可以用於分散式訓練(作為引數伺服器),或者在實際應用中提供策略。

伯克利AI分散式框架Ray,相容TensorFlow、PyTorch與MXNet

圖左:actor 為客戶端程式提供預測/操作。圖右:多個引數伺服器 actor 使用多個工作程式執行分散式訓練。


引數伺服器示例

一個引數伺服器可以作為一個 Ray actor 按如下程式碼實現:

@ray.remote
class ParameterServer(object):
 def __init__(self, keys, values):
     # These values will be mutated, so we must create a local copy.
     values = [value.copy() for value in values]
     self.parameters = dict(zip(keys, values))
 def get(self, keys):
     return [self.parameters[key] for key in keys]
 def update(self, keys, values):
     # This update function adds to the existing values, but the update
     # function can be defined arbitrarily.
     for key, value in zip(keys, values):
         self.parameters[key] += value

這裡有更完整的示例:http://ray.readthedocs.io/en/latest/example-parameter-server.html

執行以下程式碼初始化引數伺服器:

parameter_server = ParameterServer.remote(initial_keys, initial_values)

執行以下程式碼,建立 4 個長時間執行的持續恢復和更新引數的工作程式:

@ray.remote
def worker_task(parameter_server):
 while True:
     keys = ['key1', 'key2', 'key3']
     # Get the latest parameters.
     values = ray.get(parameter_server.get.remote(keys))
     # Compute some parameter updates.
     updates = …
     # Update the parameters.
     parameter_server.update.remote(keys, updates)
# Start 4 long-running tasks.
for _ in range(4):
 worker_task.remote(parameter_server)

Ray 高階庫

Ray RLib 是一個可擴充套件的強化學習庫,其建立的目的是在多個機器上執行,可以通過示例訓練指令碼或者 Python API 進行使用。目前它已有的實現為:

  • A3C
  • DQN
  • Evolution Strategies
  • PPO

UC Berkeley 的開發者在未來將繼續新增更多的演算法。同時,RLib 和 OpenAI gym 是完全相容的。

Ray.tune 是一個高效的分散式超引數搜尋庫。它提供了一個 Python API 以執行深度學習、強化學習和其它計算密集型任務。以下是一個虛構示例的程式碼:

from ray.tune import register_trainable, grid_search, run_experiments
# The function to optimize. The hyperparameters are in the config
# argument.
def my_func(config, reporter):
 import time, numpy as np
 i = 0
 while True:
     reporter(timesteps_total=i, mean_accuracy=(i ** config['alpha']))
     i += config['beta']
     time.sleep(0.01)
register_trainable('my_func', my_func)
run_experiments({
 'my_experiment': {
     'run': 'my_func',
     'resources': {'cpu': 1, 'gpu': 0},
     'stop': {'mean_accuracy': 100},
     'config': {
         'alpha': grid_search([0.2, 0.4, 0.6]),
         'beta': grid_search([1, 2]),
     },
 }
})

可以使用 TensorBoard 和 rllab's VisKit 等工具對執行狀態的結果進行實時視覺化(或者直接閱讀 JSON 日誌)。Ray.tune 支援網格搜尋、隨機搜尋和更復雜的早停演算法,如 HyperBand。

原文地址:http://bair.berkeley.edu/blog/2018/01/09/ray/

相關文章