Execution Engine
MXNet 的執行引擎不僅僅是爲了深度學習和其他任何特定的領域問題. 相反地, 它設計用來解決通用問題: 根據依賴關係來執行一系列的功能操作. 有依賴關係的任意兩個功能需要被序列化. 沒有依賴的功能 可以 併發執行來提升系統性能. 也可以參考 Note on Dependency Engine.
Interface
執行引擎的核心接口如下:
virtual void PushSync(Fn exec_fun, Context exec_ctx,
std::vector<VarHandle> const& const_vars,
std::vector<VarHandle> const& mutate_vars) = 0;
這個 API 用戶將一個 函數 (exec_fun)連同它的上下文信息以及依賴關係 push 到執行引擎. exec_ctx 是 exec_fun 執行的上下文環境. const_vars 代表的是函數只有讀取權限的變量, mutate_vars 表示的是函數可以修改的變量.
先不考慮具體的細節, 執行引擎保證下面的規則:
任意兩個會修改同一個變量的函數,會根據它們 push 到引擎的順序進行序列化.
Function
執行引擎需要的函數類型按照下面方式來定義:
using Fn = std::function<void(RunContext)>;
RunContext 包含了引擎確定的運行時信息:
struct RunContext {
// stream pointer which could be safely cast to
// cudaStream_t* type
void *stream;
};
用戶也可以使用 mxnet::engine::DAGEngine::Fn 定義作爲第二種選擇, 它們的類型是一樣的.
所有的函數都會被 engine 內部的線程來執行. 在這個模型中, 我們不鼓勵用戶將 阻塞 (blocking) 函數 push 到引擎 ( 通常是處理 I/O 任務的函數, 比如讀取硬盤, web 服務, UI, 等等). 因爲阻塞函數會佔用執行線程, 同時降低了這個系統的吞吐量. 這種情況下, 我們提供了另外的一種 asynchronous 函數類型:
using Callback = std::function<void()>;
using AsyncFn = std::function<void(RunContext, Callback)>;
在 AsyncFn 函數中, 用戶可以將重要的計算交由自己的線程來執行, 同時不用等待函數執行結束. 除非異步函數的 Callback 被調用,
否則引擎不會考慮函數是否已經結束的事情.
Context
用戶可以指定函數執行的需要的 Context. 這個 Context 一般包括函數是否執行在
CPU 或者 GPU 上, 如果是 GPU, 那麼具體是哪個 GPU. Context 和 RunContext 是不一樣的. Context 包括設備類型
(gpu/cpu) 和設備 id, 而 RunContext 包含的是只有在運行時纔可以確定的信息, 比如說函數要在哪個 stream 上執行.
VarHandle
VarHandle 是 用來指定函數的依賴關係的. MXNet 執行引擎的設計目的是爲了接口和 MXNet 的其他模塊解耦合. 所以VarHandle 類似引擎爲用戶提供用來代表函數需要或者會修改的外部資源的一個令牌.
它被設計成輕量級的, 所以創建, 刪除或者拷貝一個變量只需要一點點開銷. 對於正在推送到引擎的函數, 用戶需要在 const_vars vector 裏指定需要的不可變變量, 在mutate_vars vector
裏指定會被修改的變量. 執行引擎解析函數之間的依賴關係唯一的規則是:
*任意兩個會修改同一個變量的函數,會根據它們 push 到引擎的順序進行序列化.
舉個例子, 如果 Fn1和Fn2都要修改 V2,
那麼如果 Fn2 後於 Fn1 push
到引擎, 那麼引擎會保證 Fn2 會在 Fn1 之後執行.
另一方面, 如果 Fn1 和 Fn2 都使用但不修改 V2,
那麼它們具體的執行順序是任意的.
這種設計方式可以允許引擎可以調度 狀態改變(state-mutating) 操作. 比如說, 在 DNN 中權重更新的函數可以用 += 操作來原地更新權重, 而不是每次都產生一個新的權重數組.
如果要創建一個變量, 使用 NewVar() API. 如果要刪除一個變量, 使用 PushDelete API.
Push & Wait
所有的 Push API 是異步的. API 調用會在調用之後馬上返回, 而不管 Fn 是否執行完與否.
這允許執行引擎可以在用戶的線程推送函數到引擎的時候馬上開始計算. 所有的 Push API 不是 thread-safe 的. 具體來說, 就是應該只有一個線程來調用 API.
如果你想要等待一個具體的 Fn 完成, 需要包含一個 callback 函數, 然後在你的 Fn 的最後調用前面的回調函數.
如果你要等待所有的對一個確定的變量的使用 (讀取/修改) 的 Fn, 那麼你應該調用 WaitForVar(var) API.
如果你要等待所有推送到引擎的 Fn 都結束, 那麼需要調用 WaitForAll() API.
Save Object Creation Cost
有些情況下, 你需要推送幾個函數到引擎很多次. 如果這些函數的計算是輕量級的, 那麼拷貝 Lambda 表達式和創建讀/寫變量的列表的開銷的比例就會相當高. 我們提供了一個 API 來提前創建 OprHandle:
virtual OprHandle NewOperator(AsyncFn fn,
std::vector<VarHandle> const& const_vars,
std::vector<VarHandle> const& mutate_vars) = 0;
這樣你就可以一直推送 OprHandle, 而不用每次都要創建:
virtual void Push(OprHandle op, Context exec_ctx) = 0;
可以通過調用 DeleteOperator(OprHandle op) 來刪除它. 不過一定要確保它們已經執行完了.
API Reference
.. doxygenclass:: mxnet::Engine
:members:原文鏈接：https://github.com/dmlc/mxnet/blob/master/docs/zh/architecture/dep_engine.md