柚子快報(bào)激活碼778899分享：mxnet系統(tǒng)架構(gòu)
mxnet系統(tǒng)架構(gòu)
MXNet 是一個(gè)高性能、靈活的深度學(xué)習(xí)框架，最早由李沐（Mu Li）等人開(kāi)發(fā)，并且得到了 Amazon 的支持。它支持多種語(yǔ)言（包括 Python、Scala、C++、R、Julia、Perl 等），并以其靈活的編程模型、高效的自動(dòng)微分和分布式訓(xùn)練能力而聞名。
原文鏈接
MXNet 系統(tǒng)架構(gòu)（翻譯）
這張圖展示了 MXNet 的主要模塊和組件，以及它們之間的交互。這些模塊是：
Runtime Dependency Engine: 根據(jù)運(yùn)算之間讀寫的依賴關(guān)系，調(diào)度并執(zhí)行它們。Storage Allocator: 高效地分配和回收主機(jī)（CPU）和設(shè)備（GPU）的內(nèi)存。Resource Manager: 管理全局資源，比如隨機(jī)數(shù)生成器和臨時(shí)空間。NDArray: 動(dòng)態(tài)的異步的 n-維數(shù)組，為 MXNet 提供靈活的命令式編程。Symbolic Execution: 靜態(tài)符號(hào)圖執(zhí)行器，提供高效的符號(hào)圖執(zhí)行和優(yōu)化。Operator: 定義靜態(tài)的前向和梯度計(jì)算（backprop）的運(yùn)算符。SimpleOp: 以統(tǒng)一的方式擴(kuò)展 NDArray 運(yùn)算符和符號(hào)式運(yùn)算符的運(yùn)算符。Symbol Construction: 提供創(chuàng)建計(jì)算圖（網(wǎng)絡(luò)配置）的方式。KVStore: 鍵值存儲(chǔ)接口，提供高效的參數(shù)同步機(jī)制。Data Loading (IO): 高效的數(shù)據(jù)加載和更新。
MXNet 系統(tǒng)模塊
Execution Engine (執(zhí)行引擎)
你不僅可以使用 MXNet 引擎進(jìn)行深度學(xué)習(xí)，還可以使用它來(lái)解決任何專業(yè)領(lǐng)域的問(wèn)題。MXNet 的設(shè)計(jì)目標(biāo)是解決通用問(wèn)題：根據(jù)依賴關(guān)系來(lái)執(zhí)行一系列的函數(shù)。有依賴關(guān)系的任意兩個(gè)函數(shù)應(yīng)該按順序執(zhí)行。為提升性能，沒(méi)有依賴關(guān)系的函數(shù)可以被并行執(zhí)行。關(guān)于這方面的更詳細(xì)的討論，請(qǐng)見(jiàn) notes on the dependency engine。
Interface (接口)
下面這個(gè) API 是執(zhí)行引擎的核心接口：
virtual void PushSync(Fn exec_fun, Context exec_ctx,
std::vector
std::vector
這個(gè) API 允許你將一個(gè)函數(shù)（exec_fun）和它的執(zhí)行上下文（context）以及依賴關(guān)系一起推送給執(zhí)行引擎。exec_ctx 是 exec_fun 的執(zhí)行上下文，const_vars 是函數(shù)要讀取的變量，mutate_vars 是函數(shù)要修改的變量。執(zhí)行引擎提供以下保證：
任意兩個(gè)函數(shù)，如果它們要修改同一個(gè)變量，則它們的執(zhí)行順序與它們被推送到引擎的順序是一致的。
Function (函數(shù))
引擎的函數(shù)的類型是：
using Fn = std::function
RunContext 包含了運(yùn)行時(shí)的信息，這些信息由引擎來(lái)確定。
struct RunContext {
// stream pointer which could be safely cast to
// cudaStream_t* type
void* stream
};
或者，你也可以使用 mxnet::engine::DAGEngine::Fn, 它有相同的類型定義。
所有的函數(shù)都在引擎的內(nèi)部線程中執(zhí)行。在這種模型中，把會(huì)阻塞的函數(shù)（通常是處理磁盤或網(wǎng)絡(luò)之類的 I/O 任務(wù)的操作）發(fā)送給引擎通常不是個(gè)好主意，因?yàn)樗鼤?huì)占用執(zhí)行線程，并且降低吞吐量。在這種情況下，我們提供了另一個(gè)異步的函數(shù)：
using Callback = std::function
using AsyncFn = std::function
在 AsyncFn 中，你可以把阻塞的部分傳回到你自己的線程，然后退出函數(shù)。引擎在調(diào)用 Callback 之后才會(huì)認(rèn)為 AsyncFn 函數(shù)執(zhí)行完成。
Context (上下文)
你可以指定函數(shù)執(zhí)行的上下文。這通常包括函數(shù)該在 CPU 還是 GPU 上運(yùn)行，如果指定使用 GPU，還可以指定哪個(gè) GPU。Context 和 RunContext 不同，Context 包含設(shè)備類型 (GPU/CPU) 和設(shè)備 id，而 RunContext 包含只有運(yùn)行時(shí)才能確定的信息，比如函數(shù)應(yīng)該在哪個(gè)流上執(zhí)行。
VarHandle
VarHandle 是用來(lái)指定函數(shù)之間的依賴關(guān)系的。MXNet 引擎被設(shè)計(jì)為與其他模塊之間是松耦合的，因此 VarHandle 就像一個(gè)引擎提供的標(biāo)記，你可以用它來(lái)代表函數(shù)能使用或修改的外部資源。VarHandle 很輕量，所以創(chuàng)建、刪除或復(fù)制都只有很小的開(kāi)銷。在把函數(shù)推送給引擎時(shí)，你需要在 const_vars 向量中指定函數(shù)將要使用（只讀）的變量，在 mutate_vars 向量中指定函數(shù)將要修改的變量。引擎使用以下規(guī)則來(lái)解析函數(shù)之間的依賴關(guān)系：
如果兩個(gè)函數(shù)修改至少一個(gè)共同的變量，那么它們的執(zhí)行順序和它們被推送的順序是一致的。
舉個(gè)例子，如果 Fn1 和 Fn2 都要修改 V2，并且 Fn2 是在 Fn1 之后被推送的，則執(zhí)行引擎保證 Fn2 在 Fn1 之后執(zhí)行。如果 Fn1 和 Fn2 都使用（只讀）V2，則它們的執(zhí)行順序是隨機(jī)的。
這樣的設(shè)計(jì)允許引擎最小化內(nèi)存分配的方式來(lái)調(diào)度運(yùn)算。例如，DNN 的權(quán)重更新函數(shù)可以使用 += 操作來(lái)進(jìn)行原地操作，而不是每次都生成新的數(shù)組。
你可以使用 NewVar() 來(lái)創(chuàng)建變量，用 PushDelete() 來(lái)刪除變量。
Push and Wait (推送和等待)
所有的 Push API 都是異步的，函數(shù)調(diào)用會(huì)立即返回，不管 Fn 是否執(zhí)行完成。這允許引擎在用戶線程推送函數(shù)的時(shí)候并行開(kāi)始執(zhí)行計(jì)算。Push API 不是線程安全的，在同一個(gè)時(shí)刻，只有一個(gè)線程可以調(diào)用 Push API。
如果你想要等待某個(gè)特定的 Fn 執(zhí)行完成，你可以傳入一個(gè) Callback，并且在你的 Fn 結(jié)束的時(shí)候調(diào)用它。
如果你想等待與某個(gè)變量相關(guān)的所有 Fn 都結(jié)束，可以使用 WaitForVar(var) API。
如果你想等待所有已推送的 Fn 都結(jié)束，可以使用 WaitForAll() API。
Save Object Creation Cost (減少創(chuàng)建對(duì)象的開(kāi)銷)
在某些情況下，你需要在較長(zhǎng)一段時(shí)間內(nèi)把多個(gè)函數(shù)推送到引擎。如果這些函數(shù)的計(jì)算量不大，那么復(fù)制匿名函數(shù)和創(chuàng)建變量的開(kāi)銷就會(huì)變得相對(duì)比較高。這種情況下，我們提供了 API 來(lái)提前創(chuàng)建 OprHandle：
virtual OprHandle NewOperator(AsyncFn fn,
std::vector
std::vector
你可以連續(xù)推送 OprHandle 而不用重復(fù)創(chuàng)建它們：
virtual void Push(OprHandle op, Context exec_ctx) = 0;
要?jiǎng)h除它，調(diào)用 DeleteOperator(OprHandle op) API。要保證在調(diào)用這個(gè) API 之前，運(yùn)算符已經(jīng)完成計(jì)算。
API Reference
略。
Operators (運(yùn)算符) in MXNet
在 MXNet 中，運(yùn)算符是一個(gè)類，包括了實(shí)際的計(jì)算邏輯和一些幫助系統(tǒng)進(jìn)行優(yōu)化的輔助信息，比如原地更新和自動(dòng)求導(dǎo)之類的。要理解這篇文檔余下的部分，我們建議你熟悉一下 mshadow 庫(kù)，因?yàn)樗械倪\(yùn)算符都是在系統(tǒng)運(yùn)行時(shí)提供的類似張量（tensor-like）的數(shù)據(jù)結(jié)構(gòu) mshadow::TBlob 上進(jìn)行計(jì)算。
MXNet 的運(yùn)算符接口允許你：
通過(guò)指定原地更新來(lái)減少內(nèi)存分配。對(duì) Python 接口隱藏一些內(nèi)部參數(shù)，使接口更簡(jiǎn)潔。定義輸入張量和輸出張量之間的關(guān)系，允許系統(tǒng)為你檢查它們的形狀（shape）。為執(zhí)行計(jì)算（如調(diào)用 cudnn 的函數(shù)）向系統(tǒng)請(qǐng)求額外的臨時(shí)空間。
Operator Interface (運(yùn)算符接口)
Forward 是核心的運(yùn)算符接口：
virtual void Forward(const OpContext &ctx,
const std::vector
const std::vector
const std::vector
const std::vector
OpContext 結(jié)構(gòu)體：
struct OpContext {
int is_train;
RunContext run_ctx;
std::vector
}
它描述了此運(yùn)算符是在訓(xùn)練階段還是測(cè)試階段，它應(yīng)該運(yùn)行在哪個(gè)設(shè)備上（run_ctx），以及已經(jīng)請(qǐng)求的資源（這個(gè)會(huì)在之后的章節(jié)討論）。
in_data 和 out_data 分別代表輸入和輸出張量。系統(tǒng)已經(jīng)為所有張量分配好了空間。 req 表示計(jì)算結(jié)果如何寫入 out_data。換句話說(shuō)，req.size() == out_data.size()，并且 req[i] 對(duì)應(yīng)于 out_data[i] 的寫入類型。 OpReqType 的定義為： enum OpReqType {
kNullOp,
kWriteTo,
kWriteInplace,
kAddTo
};
通常情況下，所有 out_data 的類型應(yīng)該為 kWriteTo，表明所提供的 out_data 張量是原始的內(nèi)存塊，運(yùn)算符應(yīng)當(dāng)直接向它里面寫入數(shù)據(jù)。但是在某些情況下，比如在計(jì)算梯度張量時(shí)，我們最好對(duì)結(jié)果進(jìn)行累加，而不是直接覆蓋張量原有的內(nèi)容，這樣我們就不用每次分配額外的內(nèi)存。在這種情況下，相對(duì)應(yīng)的 req 類型應(yīng)被設(shè)置為 kAddTo，表示應(yīng)當(dāng)調(diào)用 +=。 aux_states 被設(shè)計(jì)為輔助計(jì)算的張量，目前沒(méi)有用到。
除了 Foward 操作，你可以視需要選擇實(shí)現(xiàn) Backward 接口：
virtual void Backward(const OpContext &ctx,
const std::vector
const std::vector
const std::vector
const std::vector
const std::vector
const std::vector
這個(gè)接口遵循與 Forward 相同的設(shè)計(jì)原則，不同之處是接口中 out_grad, in_data 和 out_data 是給定的，需要計(jì)算 in_grad 作為結(jié)果。這里的命名規(guī)則與 Torch 類似，可以總結(jié)為下圖：
[input/output semantics figure]
某些運(yùn)算符可能可以省略某個(gè)參數(shù)：out_grad, in_data, 和 out_data。你可以用 OperatorProperty 類的 DeclareBackwardDependency 接口來(lái)指定它們的依賴關(guān)系。
Operator Property (運(yùn)算符屬性)
卷積可以有多種實(shí)現(xiàn)方式，你可能想在各種方式之間切換以實(shí)現(xiàn)最佳性能。因此，我們把運(yùn)算符的語(yǔ)義接口從他的實(shí)現(xiàn)接口（Operator 類）中剝離出來(lái)，放到 OperatorProperty 類中。OperatorProperty 接口包含：
InferShape virtual bool InferShape(std::vector
std::vector
std::vector
這個(gè)接口有兩個(gè)目的：
告訴系統(tǒng)每個(gè)輸入張量和輸出張量的形狀，以便在調(diào)用 Forward 和 Backward 之前分配空間。在執(zhí)行之前做檢查，保證沒(méi)有明顯的錯(cuò)誤。in_shape 中指定的形狀是系統(tǒng)設(shè)置的（從前一個(gè)操作的 out_shape 中得到）。當(dāng)已知的信息不足以推斷出形狀時(shí)，InferShape 返回 false，并且當(dāng)參數(shù)形狀不一致時(shí)會(huì)拋出異常。請(qǐng)求資源: 像 cudnnConvolutionForward 之類的運(yùn)算符在計(jì)算時(shí)需要一個(gè)工作區(qū)（workspace）。如果系統(tǒng)能夠管理工作區(qū)，就可以對(duì)它進(jìn)行優(yōu)化，比如重用空間等。為此，MXNet 定義了兩個(gè)接口： virtual std::vector
virtual std::vector
ResourceRequest 結(jié)構(gòu)體（在 resource.h 中）目前僅包含一個(gè)類型標(biāo)志： struct ResourceRequest {
enum Type {
kRandom, // get a mshadow::Random object
kTempSpace, // request temporary space
};
Type type;
};
如果 ForwardResource 和 BackwardResource 返回非空的數(shù)組，系統(tǒng)會(huì)通過(guò) Operator 類的 Forward 和 Backward 接口的 ctx 參數(shù)，來(lái)提供相應(yīng)的資源。大體上說(shuō)，要訪問(wèn)這些資源，使用： auto tmp_space_res = ctx.requested[kTempSpace].get_space(some_shape, some_stream);
auto rand_res = ctx.requested[kRandom].get_random(some_stream);
示例請(qǐng)見(jiàn) src/operator/cudnn_convolution-inl.h Backward dependency (反向依賴): 讓我們看一下兩個(gè)不同的運(yùn)算符的的函數(shù)簽名（為方便展示，我們給每個(gè)參數(shù)加上了名字）： void FullyConnectedForward(TBlob weight, TBlob in_data, TBlob out_data);
void FullyConnectedBackward(TBlob weight, TBlob in_data, TBlob out_grad, TBlob in_grad);
void PoolingForward(TBlob in_data, TBlob out_data);
void PoolingBackward(TBlob in_data, TBlob out_data, TBlob out_grad, TBlob in_grad);
注意 FullyConnectedForward 中的 out_data 沒(méi)有在 FullyConnectedBackward 中被用到，而 PoolingBackward 用到了 PoolingForward 的所有參數(shù)。因此對(duì)于 FullyConnectedForward，out_data 張量在使用完之后馬上可以釋放空間，因?yàn)橄鄳?yīng)的反向函數(shù)不需要它。這給系統(tǒng)提供了機(jī)會(huì)來(lái)盡早做垃圾回收。我們提供了一個(gè)接口來(lái)指定： virtual std::vector
const std::vector
const std::vector
const std::vector
參數(shù) vector 中的 int 是 ID，用來(lái)區(qū)分不同的數(shù)組。讓我們看一下這個(gè)接口是如何為 FullyConnected 和 Pooling 指定不同的依賴關(guān)系： std::vector
const std::vector
const std::vector
const std::vector
return {out_grad[0], in_data[0]}; // NOTE: out_data[0] is NOT included
}
std::vector
const std::vector
const std::vector
const std::vector
return {out_grad[0], in_data[0], out_data[0]};
}
In place Option (原地更新選項(xiàng))：為了節(jié)省更多的內(nèi)存，你可以使用原地更新（in-place updates）。它們適用于輸入張量和輸出張量有相同形狀時(shí)的元素操作（element-wise operations）。你使用以下接口來(lái)指定原地更新： virtual std::vector
const std::vector
const std::vector
return { {in_data[0], out_data[0]} };
}
virtual std::vector
const std::vector
const std::vector
const std::vector
const std::vector
return { {out_grad[0], in_grad[0]} }
}
這段代碼告訴系統(tǒng)，在 Forward 中，in_data[0] 和 out_data[0] 張量可以共用同一塊內(nèi)存空間，而在 Backward 中，out_grad[0] 和 in_grad[0] 可以共用空間。
重要: 即使你按照以上代碼指定了共享選項(xiàng)，系統(tǒng)也不保證輸入和輸出張量會(huì)共享同一塊空間。實(shí)際上，這只是給系統(tǒng)一個(gè)建議，最終還是系統(tǒng)自己來(lái)決定是否要共用空間。不管怎樣，這個(gè)決定對(duì)你來(lái)說(shuō)是透明的，所以在實(shí)現(xiàn) Forward 和 Backward 的時(shí)候，不需要考慮這些。
Expose Operator to Python (將運(yùn)算符暴露給 Python): 因?yàn)?C++ 的限制，你需要實(shí)現(xiàn)以下接口： // initial the property class from a list of key-value string pairs
virtual void Init(const vector
// return the parameters in a key-value string map
virtual map
// return the name of arguments (for generating signature in python)
virtual vector
// return the name of output values
virtual vector
// return the name of auxiliary states
virtual vector
// return the number of output values
virtual int NumOutputs() const;
// return the number of visible outputs
virtual int NumVisibleOutputs() const;
從 Operator Property 創(chuàng)建 Operator
OperatorProperty 包含了 Operator 的所有的語(yǔ)義屬性。它也負(fù)責(zé)為實(shí)際計(jì)算創(chuàng)建 Operator。
創(chuàng)建 Operator
實(shí)現(xiàn) OperatorProperty 中的如下這個(gè)接口：
virtual Operator* CreateOperator(Context ctx) const = 0;
例如：
class ConvolutionOp {
public:
void Forward( ... ) { ... }
void Backward( ... ) { ... }
};
class ConvolutionOpProperty : public OperatorProperty {
public:
Operator* CreateOperator(Context ctx) const {
return new ConvolutionOp;
}
};
Operator 的參數(shù)化
當(dāng)你實(shí)現(xiàn)一個(gè)卷積運(yùn)算符時(shí)，你需要知道核的大小 (kernel size)，步長(zhǎng)的大小 (stride size)，填充的大小 (padding size)，等等。這些參數(shù)應(yīng)當(dāng)在 Forward 和 Backward 接口被調(diào)用之前傳給 Operator。你可以定義一個(gè) ConvolutionParam 結(jié)構(gòu)，如下：
#include
struct ConvolutionParam : public dmlc::Parameter
TShape kernel, stride, pad;
uint32_t num_filter, num_group, workspace;
bool no_bias;
};
把它放在 ConvolutionOpProperty 中，并且在創(chuàng)建 Operator 時(shí)傳進(jìn)去：
class ConvolutionOp {
public:
ConvolutionOp(ConvolutionParam p): param_(p) {}
void Forward( ... ) { ... }
void Backward( ... ) { ... }
private:
ConvolutionParam param_;
};
class ConvolutionOpProperty : public OperatorProperty {
public:
void Init(const vector
// initialize param_ using kwargs
}
Operator* CreateOperator(Context ctx) const {
return new ConvolutionOp(param_);
}
private:
ConvolutionParam param_;
};
使用以下宏來(lái)把 Operator 的 Property 類和 Parameter 類注冊(cè)到 MXNet：
DMLC_REGISTER_PARAMETER(ConvolutionParam);
MXNET_REGISTER_OP_PROPERTY(Convolution, ConvolutionOpProperty);
第一個(gè)參數(shù)是名字，第二個(gè)參數(shù)是 Property 的類名。
接口總結(jié)
到目前為止，我們基本上涵蓋了定義一個(gè) Operator 的全部接口。讓我們回顧一下：
使用 Operator 接口來(lái)實(shí)現(xiàn)你的計(jì)算邏輯 (Forward 和 Backward)。使用 OperatorProperty 接口來(lái)：
向運(yùn)算符類傳遞參數(shù)（可以使用 Init 接口）使用 CreateOperator 接口來(lái)創(chuàng)建運(yùn)算符正確地實(shí)現(xiàn)描述操作符的接口，例如參數(shù)名，等等正確地實(shí)現(xiàn) InferShape 接口來(lái)設(shè)置輸出張量的形狀[可選] 如果需要額外的資源，檢查 ForwardResource 和 BackwardResource[可選] 如果 Backward 不需要用到 Forward 的所有輸入和輸出，檢查 DeclareBackwardDependency[可選] 如果支持原地更新，檢查 ForwardInplaceOption 和 BackwardInplaceOption 將 OperatorProperty 類和參數(shù)類注冊(cè)到 MXNet
統(tǒng)一 NDArray 運(yùn)算符和符號(hào)運(yùn)算符
NDArray 運(yùn)算符和符號(hào)運(yùn)算符類似，區(qū)別是在沒(méi)有完整的依賴關(guān)系圖時(shí)，有時(shí)你不能原地更新。然而，NDArray 運(yùn)算符和符號(hào)運(yùn)算符的底層邏輯是一樣的。SimpleOp 是一個(gè)新的統(tǒng)一的運(yùn)算符 API，統(tǒng)一了不同的方式。因?yàn)槎鄶?shù)數(shù)學(xué)運(yùn)算符有一個(gè)或兩個(gè)操作數(shù)，而更多個(gè)操作數(shù)使得和依賴關(guān)系相關(guān)的優(yōu)化有用，統(tǒng)一的運(yùn)算符是被專門設(shè)計(jì)用于一元和二元運(yùn)算。
讓我們考慮運(yùn)算符的基本元素。理想情況下，你只需要用函數(shù)和導(dǎo)數(shù)來(lái)描述一個(gè)運(yùn)算符。讓我們將討論限制在一元和二元運(yùn)算符。我們?cè)撊绾畏诸愃胁僮鞣?，?lái)使原地更新的可能性最大？注意你可以按照操作數(shù)的數(shù)量來(lái)對(duì)函數(shù)進(jìn)行分類。而導(dǎo)數(shù)更復(fù)雜一些。要?jiǎng)?chuàng)建一個(gè)依賴關(guān)系圖，你需要知道輸出值，輸入數(shù)據(jù)是否在之后的梯度中被用到。統(tǒng)一 API 中的梯度函數(shù)是用操作數(shù)的類型來(lái)區(qū)分的。
在我們繼續(xù)了解 SimpleOp 接口之前，我們建議你瀏覽 mshadow library guide，因?yàn)橛?jì)算是在 mshadow:TBlob 中進(jìn)行的。
在以下例子中，我們會(huì)創(chuàng)建一個(gè)實(shí)現(xiàn) smooth l1 loss 函數(shù)的運(yùn)算符，這是 l1 loss 和 l2 loss 的混合體。這個(gè)函數(shù)可以被寫成如下形式：
loss = outside_weight .* f(inside_weight .* (data - label))
grad = outside_weight .* inside_weight .* f'(inside_weight .* (data - label))
.* 代表元素乘，f 和 f’ 是 smooth l1 loss 函數(shù)，我們先假設(shè)它們?cè)?mshadow 中可以找到。乍看起來(lái)，不可能把它實(shí)現(xiàn)成一元或二元操作。但是我們有符號(hào)的自動(dòng)微分，這簡(jiǎn)化了 f 和 f’ 的損失函數(shù)。這個(gè)損失函數(shù)并不比 sin 或者 abs 函數(shù)復(fù)雜，我們可以將它實(shí)現(xiàn)為一元操作符。
SimpleOp: 統(tǒng)一的 Operator API
Define Shapes (定義形狀)
mshadow 庫(kù)需要顯式地分配內(nèi)存，因此在計(jì)算開(kāi)始前，就需要定義好所有數(shù)據(jù)的形狀。在定義函數(shù)和導(dǎo)數(shù)之前，先讓我們檢查輸入的形狀并且提供輸出的形狀。
···cpp typedef TShape (*UnaryShapeFunction)(const TShape& src,const EnvArguments& env); typedef TShape (*BinaryShapeFunction)(const TShape& const lhs, TShape& rhs, const EnvArguments& env); ···
你可以使用 mshadow::TShape 來(lái)檢查輸入形狀并且指定輸出形狀。如果你不定義這個(gè)函數(shù)，則默認(rèn)的輸出形狀和輸入形狀相同。如果是二元運(yùn)算符，默認(rèn)情況下，系統(tǒng)會(huì)檢查 lhs 和 rhs 的形狀是否相同。
你還可以用這些函數(shù)來(lái)檢查是否有額外的參數(shù)和資源?？梢詤⒖?EnvArguments 的用法。
在我們開(kāi)始 smooth l1 loss 的例子之前，我們?cè)陬^文件 smooth_l1_unary-inl.h 中定義了 XPU，值為 cpu 或者 gpu，以便于我們可以在 smooth_l1_unary.cc 和 smoth_l1_unary.cu 中使用相同的代碼。
#include
#if defined(__CUDACC__)
#define XPU gpu
#else
#define XPU cpu
#endif
在 smooth l1 loss 的例子中，因?yàn)檩斎胼敵鲇邢嗤男螤?，我們就直接用默認(rèn)行為：
inline TShape SmoothL1Shape_(const TShape& src, const EnvArguments& env) {
return TShape(src);
}
定義函數(shù)
創(chuàng)建有一個(gè)輸出 (mshadow::TBlob) 的一元或二元函數(shù)。
typedef void (*UnaryFunction)(const TBlob& src,
const EnvArguments& env,
TBlob* ret,
OpReqType req,
RunContext ctx);
typedef void (*BinaryFunction)(const TBlob& lhs,
const TBlob& rhs,
const EnvArguments& env,
TBlob* ret,
OpReqType req,
RunContext ctx);
函數(shù)按照輸入?yún)?shù)的類型區(qū)分。 RunContext ctx 包含運(yùn)行時(shí)所需要的信息。 struct RunContext {
void *stream; // the stream of the device, can be NULL or Stream* in GPU mode
template
} // namespace mxnet
從 ctx 獲取一個(gè)流的例子: mshadow::stream *s = ctx.get_stream();
OpReqType req 指明計(jì)算結(jié)果該如何寫入 ret: enum OpReqType {
kNullOp, // no operation, do not write anything
kWriteTo, // write gradient to provided space
kWriteInplace, // perform an in-place write
kAddTo // add to the provided space
};
ASSIGN_DISPATH(out, req, exp) 是 operator_util.h 中定義的一個(gè)宏，用來(lái)簡(jiǎn)化 OpReqType 的使用，它會(huì)檢查 req 并且進(jìn)行賦值。
在 smooth l1 loss 例子中，我們使用 UnaryFunction 來(lái)定義操作符的函數(shù)。
template
void SmoothL1Forward_(const TBlob& src,
const EnvArguments& env,
TBlob *ret,
OpReqType req,
RunContext ctx) {
using namespace mshadow;
using namespace mshadow::expr;
mshadow::Stream
real_t sigma2 = env.scalar * env.scalar;
MSHADOW_TYPE_SWITCH(ret->type_flag_, DType, {
mshadow::Tensor
mshadow::Tensor
ASSIGN_DISPATCH(out, req,
F
});
}
從 RunContext 獲取到 mshadow::Stream 之后，我們從 mshadow::TBlob 拿到 mshadow::Tensor。mshadow::F 是創(chuàng)建一個(gè) mshadow 表達(dá)式的快捷方式。宏 MSHADOW_TYPE_SWITCH(type, DType, …) 處理不同類型的細(xì)節(jié)，宏 ASSIGN_DISPATCH(out, req, exp) 檢查 OpReqType 并且執(zhí)行相應(yīng)動(dòng)作。 sigma2 是這個(gè)損失函數(shù)中的一個(gè)特殊的參數(shù)，我們后面會(huì)講到。
定義梯度（可選）
// depending only on out_grad
typedef void (*UnaryGradFunctionT0)(const OutputGrad& out_grad,
const EnvArguments& env,
TBlob* in_grad,
OpReqType req,
RunContext ctx);
// depending only on out_value
typedef void (*UnaryGradFunctionT1)(const OutputGrad& out_grad,
const OutputValue& out_value,
const EnvArguments& env,
TBlob* in_grad,
OpReqType req,
RunContext ctx);
// depending only on in_data
typedef void (*UnaryGradFunctionT2)(const OutputGrad& out_grad,
const Input0& in_data0,
const EnvArguments& env,
TBlob* in_grad,
OpReqType req,
RunContext ctx);
二元運(yùn)算符的梯度函數(shù)擁有相似的結(jié)構(gòu)，不同之處在于 Input, TBlob 和 OpReqType 的數(shù)量翻倍。
GradFunctionArgument
Input0, Input, OutputValue 和 OutputGrad 都和 GradFunctionArgument 有相同的結(jié)構(gòu)，定義如下：
struct GradFunctionArgument {
TBlob data;
}
在 smooth l1 loss 例子中，注意用到輸入來(lái)計(jì)算梯度的是 f’(x)，所以 UnaryGradFunctionT2 是適用的。我們還需要用 out_grade 乘以結(jié)果的 in_grad 來(lái)用鏈?zhǔn)椒▌t計(jì)算梯度。
template
void SmoothL1BackwardUseIn_(const OutputGrad& out_grad,
const Input0& in_data0,
const EnvArguments& env,
TBlob *in_grad,
OpReqType req,
RunContext ctx) {
using namespace mshadow;
using namespace mshadow::expr;
mshadow::Stream
real_t sigma2 = env.scalar * env.scalar;
MSHADOW_TYPE_SWITCH(in_grad->type_flag_, DType, {
mshadow::Tensor
mshadow::Tensor
mshadow::Tensor
ASSIGN_DISPATCH(igrad, req,
ograd * F
});
}
把 SimpleOp 注冊(cè)到 MXNet
在創(chuàng)建 shape, function 和 gradient 之后，要把它們放到 NDArray 運(yùn)算符和符號(hào)運(yùn)算符中?？梢允褂?operator_util.h 中定義的宏來(lái)簡(jiǎn)化這個(gè)過(guò)程。
MXNET_REGISTER_SIMPLE_OP(Name, DEV)
.set_shape_function(Shape)
.set_function(DEV::kDevMask, Function
.set_gradient(DEV::kDevMask, Gradient
.describe("description");
SimpleOpInplaceOption 的定義：
enum SimpleOpInplaceOption {
kNoInplace, // do not allow inplace in arguments
kInplaceInOut, // allow inplace in with out (unary)
kInplaceOutIn, // allow inplace out_grad with in_grad (unary)
kInplaceLhsOut, // allow inplace left operand with out (binary)
kInplaceOutLhs // allow inplace out_grad with lhs_grad (binary)
};
在我們的例子中，我們的梯度依賴于函數(shù)的輸入，因此函數(shù)不能原地更新數(shù)據(jù)。輸出的梯度在梯度計(jì)算后就不被用到了，因此梯度可以原地更新。
MXNET_REGISTER_SIMPLE_OP(smooth_l1, XPU)
.set_function(XPU::kDevMask, SmoothL1Forward_
.set_gradient(XPU::kDevMask, SmoothL1BackwardUseIn_
.set_enable_scalar(true)
.describe("Calculate Smooth L1 Loss(lhs, scalar)");
不要忘了之前的討論，在沒(méi)有用 set_shape_function 來(lái)設(shè)置 shape 的時(shí)候，默認(rèn)會(huì)強(qiáng)制輸出的 shape 和輸入的 shape 一致。我們后面會(huì)討論 set_enable_scalar。
NDArray 運(yùn)算符總結(jié)
創(chuàng)建一個(gè) shape 函數(shù)用來(lái)決定輸出 shape創(chuàng)建一個(gè)函數(shù)作為前向過(guò)程，選擇一個(gè)合適的函數(shù)類型創(chuàng)建一個(gè)梯度作為反向過(guò)程，選擇一個(gè)合適的梯度類型注冊(cè)運(yùn)算符
SimpleOp 的其他信息
在 EnvArguments 上使用 SimpleOp
一些操作需要一個(gè)標(biāo)量作為輸入，比如一個(gè)梯度標(biāo)量，一組用來(lái)控制算法行為的關(guān)鍵字參數(shù)，或者一個(gè)用來(lái)加速計(jì)算的臨時(shí)空間。EnvArguments 提供額外的參數(shù)和資源來(lái)使計(jì)算更易擴(kuò)展和更高效。
struct EnvArguments {
real_t scalar; // scalar argument, if enabled
std::vector
std::vector
};
要啟用這些額外的功能，需要更多的注冊(cè)參數(shù)。為避免參數(shù)的混淆，scalar 和 kwargs 不能同時(shí)使用。要啟用 scalar，在注冊(cè)時(shí)使用 set_enable_scalar(bool enable_scalar)。之后，在前向函數(shù)和梯度中，可以用函數(shù)測(cè)參數(shù) EnvArguments env 中的 env.scalar 來(lái)訪問(wèn) scalar。
要啟用 kwargs，在注冊(cè)時(shí)使用 set_enable_kwargs(bool enable_kwargs)。在前向函數(shù)和反向梯度時(shí)，額外的參數(shù)包含在 env.kwargs 中，被定義為 std::vectorstd::string。使用 DMLC 參數(shù)接口來(lái)簡(jiǎn)化關(guān)鍵字參數(shù)的解析。更多細(xì)節(jié)請(qǐng)參考 guide on parameter structure。
mshadow::Random 或者臨時(shí)內(nèi)存空間之類的額外資源可以通過(guò) EnvArguments.resoure 來(lái)請(qǐng)求和訪問(wèn)。注冊(cè)方式為 set_resource_request(ResourceRequest req) 或者 set_resource_request(const std::vector)，其中 mxnet::ResourceRequest 被定義為：
struct ResourceRequest {
enum Type { // Resource type, indicating what the pointer type is
kRandom, // mshadow::Random object
kTempSpace // A dynamic temp space that can be arbitrary size
};
Type type; // type of resources
};
相關(guān)例子請(qǐng)參見(jiàn) src/operator/loss_binary_op-inl.h。
在我們的 smooth l1 loss 例子中，需要有一個(gè)標(biāo)量輸入來(lái)標(biāo)記損失函數(shù)的折點(diǎn)。因此，在注冊(cè)時(shí)，我們使用 set_enable_scalar(true)，并且在函數(shù)和梯度中使用 env.scalar。
實(shí)現(xiàn)一個(gè)張量運(yùn)算 (Tensor Operation)
因?yàn)槭褂?mshadow 庫(kù)來(lái)進(jìn)行計(jì)算，有時(shí)候沒(méi)有我們用到的函數(shù)，我們可以在運(yùn)算符中實(shí)現(xiàn)張量運(yùn)算。如果你把函數(shù)定義為元素 (element-wise) 運(yùn)算，那么你可以把它實(shí)現(xiàn)成一個(gè) mxnet::op::mshadow_op。src/operator/mshadow_op.h 中有許多 mshadow_op 的例子。 mshadow_op 是表達(dá)式映射器。它們處理函數(shù)的標(biāo)量形式。細(xì)節(jié)請(qǐng)見(jiàn) mshadow expression API guide。
如果一個(gè)運(yùn)算不能用元素 (element-wise) 方式實(shí)現(xiàn)，比如 softmax 損失函數(shù)和梯度，那么你就需要實(shí)現(xiàn)一個(gè)新的張量運(yùn)算。你需要直接創(chuàng)建 mshadow 函數(shù)和 mshadow::cuda 函數(shù)。更多例子請(qǐng)見(jiàn) src/ooperator/roi_pooling.cc。
在我們的 smooth l1 loss 例子中，我們創(chuàng)建了兩個(gè)映射器，分別是標(biāo)量情形下的 smooth l1 loss 和 gradient。
namespace mshadow_op {
struct smooth_l1_loss {
// a is x, b is sigma2
MSHADOW_XINLINE static real_t Map(real_t a, real_t b) {
if (a > 1.0f / b) {
return a - 0.5f / b;
} else if (a < -1.0f / b) {
return -a - 0.5f / b;
} else {
return 0.5f * a * a * b;
}
}
};
}
梯度 (gradient) 與之相似，可以在 src/operator/smooth_l1_unary-inl.h 中找到。
兩個(gè)以上操作數(shù)
新的統(tǒng)一 API 被設(shè)計(jì)為完成運(yùn)算符的基本功能。對(duì)于有兩個(gè)以上輸入、或一個(gè)以上輸出、或是需要更多特性的運(yùn)算符，請(qǐng)見(jiàn)原始的 Operator API。
其他文章可參考
https://mli.github.io/2015/12/03/mxnet-overview/https://blog.51cto.com/u_16213409/7336190https://blog.csdn.net/huareal/article/details/72589863
柚子快報(bào)激活碼778899分享：mxnet系統(tǒng)架構(gòu)
推薦閱讀
本文內(nèi)容根據(jù)網(wǎng)絡(luò)資料整理，出于傳遞更多信息之目的，不代表金鑰匙跨境贊同其觀點(diǎn)和立場(chǎng)。
轉(zhuǎn)載請(qǐng)注明，如有侵權(quán)，聯(lián)系刪除。
eBay跨境俠
