Apache MXNet - 系統元件

上一頁
下一頁

這裡詳細解釋了 Apache MXNet 中的系統元件。首先，我們將學習 MXNet 中的執行引擎。

執行引擎

Apache MXNet 的執行引擎非常通用。我們可以將其用於深度學習以及任何特定領域的難題：按照其依賴關係執行一堆函式。它的設計方式是，具有依賴關係的函式被序列化，而沒有依賴關係的函式可以並行執行。

核心介面

下面給出的 API 是 Apache MXNet 執行引擎的核心介面：

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 − MXNet 的核心介面 API 允許我們將名為 exec_fun 的函式及其上下文資訊和依賴項推送到執行引擎。

  • exec_ctx − 上述函式 exec_fun 應該在其執行的上下文資訊。

  • const_vars − 這些是函式從中讀取的變數。

  • mutate_vars − 這些是要修改的變數。

執行引擎向其使用者保證，任何兩個修改公共變數的函式的執行在其推送順序中是序列化的。

函式

以下是 Apache MXNet 執行引擎的函式型別：

using Fn = std::function<void(RunContext)>;

在上面的函式中，RunContext 包含執行時資訊。執行時資訊應由執行引擎確定。RunContext 的語法如下：

struct RunContext {
   // stream pointer which could be safely cast to
   // cudaStream_t* type
   void *stream;
};

以下是一些關於執行引擎函式的重要說明：

  • 所有函式都由 MXNet 執行引擎的內部執行緒執行。

  • 將阻塞函式推送到執行引擎不是一個好主意，因為這樣會佔用執行執行緒，並降低總吞吐量。

為此，MXNet 提供了另一個非同步函式，如下所示：

using Callback = std::function<void()>;
using AsyncFn = std::function<void(RunContext, Callback)>;

  • 在這個 AsyncFn 函式中，我們可以傳遞執行緒的繁重部分，但是直到我們呼叫 callback 函式，執行引擎才認為該函式已完成。

上下文

在 Context 中，我們可以指定要在其中執行函式的上下文。這通常包括以下內容：

  • 函式是否應該在 CPU 或 GPU 上執行。

  • 如果我們在 Context 中指定 GPU，則使用哪個 GPU。

  • Context 和 RunContext 之間存在巨大差異。Context 包含裝置型別和裝置 ID，而 RunContext 包含只有在執行時才能確定的資訊。

VarHandle

VarHandle 用於指定函式的依賴關係，它就像一個令牌（特別是執行引擎提供的令牌），我們可以用它來表示函式可以修改或使用的外部資源。

但是問題出現了，為什麼我們需要使用 VarHandle？這是因為 Apache MXNet 引擎的設計與其他 MXNet 模組解耦。

以下是一些關於 VarHandle 的重要說明：

  • 它很輕量級，因此建立、刪除或複製變數幾乎不會產生操作成本。

  • 我們需要指定不可變變數，即將在 const_vars 中使用的變數。

  • 我們需要指定可變變數，即將在 mutate_vars 中修改的變數。

  • 執行引擎用於解決函式之間依賴關係的規則是，當其中一個函式修改至少一個公共變數時，任何兩個函式的執行在其推送順序中是序列化的。

  • 要建立新變數，可以使用 NewVar() API。

  • 要刪除變數，可以使用 PushDelete API。

讓我們透過一個簡單的例子來了解它的工作原理：

假設我們有兩個函式 F1 和 F2，它們都修改變數 V2。在這種情況下，如果 F2 在 F1 之後被推送，則保證 F2 在 F1 之後執行。另一方面，如果 F1 和 F2 都使用 V2，則它們的實際執行順序可能是隨機的。

Push 和 Wait

Push 和 wait 是執行引擎的另外兩個有用的 API。

以下是 Push API 的兩個重要特性：

  • 所有 Push API 都是非同步的，這意味著 API 呼叫會立即返回，而不管推送的函式是否已完成。

  • Push API 不是執行緒安全的，這意味著一次只有一個執行緒應該進行引擎 API 呼叫。

現在如果我們談論 Wait API，以下幾點代表它：

  • 如果使用者想要等待特定函式完成，他/她應該在閉包中包含一個回撥函式。包含後，在函式結束時呼叫該函式。

  • 另一方面，如果使用者想要等待涉及某個變數的所有函式完成，他/她應該使用 WaitForVar(var) API。

  • 如果有人想等待所有推送的函式完成，則使用 WaitForAll() API。

  • 用於指定函式的依賴關係，就像一個令牌。

運算元

Apache MXNet 中的運算元是一個包含實際計算邏輯以及輔助資訊的類，並幫助系統執行最佳化。

運算元介面

Forward 是核心運算元介面，其語法如下：

virtual void Forward(const OpContext &ctx,
const std::vector<TBlob> &in_data,
const std::vector<OpReqType> &req,
const std::vector<TBlob> &out_data,
const std::vector<TBlob> &aux_states) = 0;

在 Forward() 中定義的 OpContext 的結構如下：

struct OpContext {
   int is_train;
   RunContext run_ctx;
   std::vector<Resource> requested;
}

OpContext 描述了運算元的狀態（無論是在訓練階段還是測試階段），運算元應該在其執行的裝置以及請求的資源。執行引擎的另外兩個有用的 API。

從上面的 Forward 核心介面，我們可以理解請求的資源如下：

  • in_data 和 out_data 表示輸入和輸出張量。

  • req 表示計算結果如何寫入 out_data。

OpReqType 可以定義為：

enum OpReqType {
   kNullOp,
   kWriteTo,
   kWriteInplace,
   kAddTo
};

與 Forward 運算元一樣，我們可以選擇性地實現 Backward 介面，如下所示：

virtual void Backward(const OpContext &ctx,
const std::vector<TBlob> &out_grad,
const std::vector<TBlob> &in_data,
const std::vector<TBlob> &out_data,
const std::vector<OpReqType> &req,
const std::vector<TBlob> &in_grad,
const std::vector<TBlob> &aux_states);

各種任務

Operator 介面允許使用者執行以下任務：

  • 使用者可以指定就地更新，並降低記憶體分配成本。

  • 為了使其更清晰，使用者可以隱藏 Python 中的一些內部引數。

  • 使用者可以定義張量和輸出張量之間的關係。

  • 要執行計算，使用者可以從系統獲取額外的臨時空間。

運算元屬性

我們知道在卷積神經網路 (CNN) 中，一個卷積有幾種實現方式。為了從這些實現中獲得最佳效能，我們可能希望在這幾種卷積之間切換。

這就是 Apache MXNet 將運算元語義介面與實現介面分開的原因。這種分離以 OperatorProperty 類的形式完成，該類包含以下內容：

InferShape − InferShape 介面有兩個目的，如下所示：

  • 第一個目的是告訴系統每個輸入和輸出張量的尺寸，以便在 Forward 和 Backward 呼叫之前分配空間。

  • 第二個目的是執行大小檢查，以確保在執行之前沒有錯誤。

語法如下：

virtual bool InferShape(mxnet::ShapeVector *in_shape,
mxnet::ShapeVector *out_shape,
mxnet::ShapeVector *aux_shape) const = 0;

請求資源 − 如果您的系統可以管理諸如 cudnnConvolutionForward 之類的操作的計算工作區呢？您的系統可以執行諸如重用空間之類的最佳化等等。在這裡，MXNet 可以藉助以下兩個介面輕鬆實現這一點：

virtual std::vector<ResourceRequest> ForwardResource(
   const mxnet::ShapeVector &in_shape) const;
virtual std::vector<ResourceRequest> BackwardResource(
   const mxnet::ShapeVector &in_shape) const;

但是，如果 ForwardResource 和 BackwardResource 返回非空陣列會怎麼樣？在這種情況下，系統透過 Operator 的 Forward 和 Backward 介面中的 ctx 引數提供相應的資源。

反向依賴 − Apache MXNet 有以下兩種不同的運算元簽名來處理反向依賴：

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 未被 FullyConnectedBackward 使用，並且

  • PoolingBackward 需要 PoolingForward 的所有引數。

這就是為什麼對於 FullyConnectedForward，一旦消耗了 out_data 張量，就可以安全地釋放它，因為反向函式不需要它。藉助這個系統，可以儘早收集一些張量作為垃圾。

就地選項 − Apache MXNet 為使用者提供了另一個介面來節省記憶體分配成本。該介面適用於輸入和輸出張量具有相同形狀的逐元素操作。

以下是指定就地更新的語法：

建立運算元的示例

藉助 OperatorProperty，我們可以建立一個運算元。為此，請按照以下步驟操作：

virtual std::vector<std::pair<int, void*>> ElewiseOpProperty::ForwardInplaceOption(
   const std::vector<int> &in_data,
   const std::vector<void*> &out_data) 
const {
   return { {in_data[0], out_data[0]} };
}
virtual std::vector<std::pair<int, void*>> ElewiseOpProperty::BackwardInplaceOption(
   const std::vector<int> &out_grad,
   const std::vector<int> &in_data,
   const std::vector<int> &out_data,
   const std::vector<void*> &in_grad) 
const {
   return { {out_grad[0], in_grad[0]} }
}

步驟 1

建立運算元

首先在 OperatorProperty 中實現以下介面：

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;
      }
};

步驟 2

引數化運算元

如果您要實現卷積運算元，則必須知道核心大小、步幅大小、填充大小等。為什麼？因為這些引數應該在呼叫任何 Forward 或 backward 介面之前傳遞給運算元。

為此，我們需要定義如下 ConvolutionParam 結構：

#include <dmlc/parameter.h>
struct ConvolutionParam : public dmlc::Parameter<ConvolutionParam> {
   mxnet::TShape kernel, stride, pad;
   uint32_t num_filter, num_group, workspace;
   bool no_bias;
};

現在，我們需要將其放入 ConvolutionOpProperty 並將其傳遞給運算元，如下所示：

class ConvolutionOp {
   public:
      ConvolutionOp(ConvolutionParam p): param_(p) {}
      void Forward( ... ) { ... }
      void Backward( ... ) { ... }
   private:
      ConvolutionParam param_;
};
class ConvolutionOpProperty : public OperatorProperty {
   public:
      void Init(const vector<pair<string, string>& kwargs) {
         // initialize param_ using kwargs
      }
      Operator* CreateOperator(Context ctx) const {
         return new ConvolutionOp(param_);
      }
   private:
      ConvolutionParam param_;
};

步驟 3

將運算元屬性類和引數類註冊到 Apache MXNet

最後，我們需要將運算元屬性類和引數類註冊到 MXNet。這可以使用以下宏來完成：

DMLC_REGISTER_PARAMETER(ConvolutionParam);
MXNET_REGISTER_OP_PROPERTY(Convolution, ConvolutionOpProperty);

在上面的宏中，第一個引數是名稱字串，第二個是屬性類名稱。

列印頁面
上一頁
下一頁
廣告