↑ 點擊藍字?關(guān)注極市平臺

作者丨Gemfield@@知乎

來源丨h(huán)ttps://zhuanlan.zhihu.com/p/61765561、https://zhuanlan.zhihu.com/p/65822256

編輯丨極市平臺

本文僅用于學術(shù)分享，著作權(quán)歸作者所有。如有侵權(quán)，請聯(lián)系后臺作刪文處理。

極市導(dǎo)讀

本文以torch/csrc/autograd下的代碼為基礎(chǔ)，深入講解了PyTorch的動態(tài)圖系統(tǒng)，可以說是目前關(guān)于PyTorch動態(tài)圖的最詳盡的文章之一。>>加入極市CV技術(shù)交流群，走在計算機視覺的最前沿

PyTorch的動態(tài)圖（上）

背景

PyTorch的動態(tài)圖框架主要是由torch/csrc/autograd下的代碼實現(xiàn)的。這個目錄下定義了3個主要的基類：Variable、Function、Engine，這三個基類及其繼承體系共同構(gòu)成了PyTorch動態(tài)圖的根基。

為什么叫作動態(tài)圖呢？圖容易理解，F(xiàn)unction是nodes/vertices，(Function, input_nr)是edges。那么動態(tài)體現(xiàn)在什么地方呢？每一次前向時構(gòu)建graph，反向時銷毀。本文就以torch/csrc/autograd/下的代碼為基礎(chǔ)，深入講解PyTorch的動態(tài)圖系統(tǒng)——這也可能是互聯(lián)網(wǎng)上關(guān)于PyTorch動態(tài)圖最詳盡的文章了。

在專欄文章《PyTorch的初始化》（https://zhuanlan.zhihu.com/p/57571317）中，gemfield描述了PyTorch的初始化流程，在文末提到了THPAutograd_initFunctions()調(diào)用：“最后的THPAutograd_initFunctions()則是初始化了torch的自動微分系統(tǒng)，這是PyTorch動態(tài)圖框架的基礎(chǔ)”。而本文將以THPAutograd_initFunctions開始，帶你走入到PyTorch的動態(tài)圖世界中。首先為上篇，主要介紹Function、Variable、Engine的類的繼承體系。

autograd初始化

THPAutograd_initFunctions這個函數(shù)實現(xiàn)如下：

void?THPAutograd_initFunctions()
{
??THPObjectPtr?module(PyModule_New("torch._C._functions"));
??......
??generated::initialize_autogenerated_functions();
??auto?c_module?=?THPObjectPtr(PyImport_ImportModule("torch._C"));
}

用來初始化cpp_function_types表，這個表維護了從cpp類型的函數(shù)到python類型的映射：

static std::unordered_map cpp_function_types

這個表里存放的都是和autograd相關(guān)的函數(shù)的映射關(guān)系，起什么作用呢？比如我在python中print一個Variable的grad_fn：

>>> gemfield = torch.empty([2,2],requires_grad=True)
>>> syszux = gemfield * gemfield
>>> syszux.grad_fn

grad_fn是一個Function的實例，我們在C++中定義了那么多反向函數(shù)（參考下文），但是怎么在python中訪問呢？就靠上面這個表的映射。實際上，cpp_function_types這個映射表就是為了在python中打印grad_fn服務(wù)的。

Variable

參考：https://zhuanlan.zhihu.com/p/64135058

以下面的代碼片段作為例子：

gemfield = torch.ones(2, 2, requires_grad=True)
syszux = gemfield + 2
civilnet = syszux * syszux * 3
gemfieldout = civilnet.mean()
gemfieldout.backward()

需要指出的是，動態(tài)圖是在前向的時候建立起來的。gemfieldout作為前向的最終輸出，在反向傳播的時候，卻是計算的最初輸入—在動態(tài)圖中，我們稱之為root。在下文介紹Engine的時候，你就會看到，我們會使用gemfieldout這個root來構(gòu)建GraphRoot實例，以此作為Graph的輸入。

Function

在開始介紹Function之前，還是以上面的代碼為例，在一次前向的過程中，我們會創(chuàng)建出如下的Variable和Function實例：

#Variable實例
gemfield --> grad_fn_ (Function實例）= None
         --> grad_accumulator_ (Function實例）= AccumulateGrad實例0x55ca7f304500
         --> output_nr_ = 0

#Function實例， 0x55ca7f872e90
AddBackward0實例 --> sequence_nr_ (uint64_t) = 0
            --> next_edges_ (edge_list) --> std::vector = [(AccumulateGrad實例, 0)，(0, 0)]
            --> input_metadata_ --> [(type, shape, device)...] = [(CPUFloatType, [2, 2],cpu])]
            --> alpha (Scalar) = 1
            --> apply() --> 使用 AddBackward0 的apply

#Variable實例
syszux --> grad_fn_ (Function實例）= AddBackward0實例0x55ca7f872e90
       --> output_nr_ = 0

#Function實例， 0x55ca7ebba2a0
MulBackward0 --> sequence_nr_ (uint64_t) = 1
            --> next_edges_ (edge_list) = [(AddBackward0實例0x55ca7f872e90,0),(AddBackward0實例0x55ca7f872e90,0)]
            --> input_metadata_ --> [(type, shape, device)...] = [(CPUFloatType, [2, 2],cpu])]
            --> alpha (Scalar) = 1
            --> apply() --> 使用 MulBackward0 的apply

# #Variable實例，syszux * syszux得到的tmp
tmp --> grad_fn_ (Function實例）= MulBackward0實例0x55ca7ebba2a0
    --> output_nr_ = 0

#Function實例，0x55ca7fada2f0
MulBackward0 --> sequence_nr_ (uint64_t) = 2 （每個線程內(nèi)自增）
            --> next_edges_ (edge_list) = [(MulBackward0實例0x55ca7ebba2a0,0),(0,0)]
            --> input_metadata_ --> [(type, shape, device)...] = [(CPUFloatType, [2, 2],cpu])]
            --> self_ (SavedVariable) = tmp的淺拷貝
            --> other_ (SavedVariable) = 3的淺拷貝
            --> apply() --> 使用 MulBackward0 的apply

#Variable實例
civilnet --> grad_fn_ (Function實例）= MulBackward0實例0x55ca7fada2f0                                          -

#Function實例，0x55ca7eb358b0
MeanBackward0 --> sequence_nr_ (uint64_t) = 3 （每個線程內(nèi)自增）
              --> next_edges_ (edge_list) = [(MulBackward0實例0x55ca7fada2f0,0)]
              --> input_metadata_ --> [(type, shape, device)...] = [(CPUFloatType|[]|cpu])]
              --> self_sizes (std::vector) = (2, 2)
              --> self_numel = 4
              --> apply() --> 使用 MulBackward0 的apply
#Variable實例
gemfieldout --> grad_fn_ (Function實例）= MeanBackward0實例0x55ca7eb358b0
            --> output_nr_ = 0

這些用于反向計算的Function實例之間通過next_edges_連接在一起，因為這些Function的實際運行都是在反向期間，因此，輸出輸出關(guān)系正好和前向期間是反過來的。它們通過next_edges_連接在一起。用一個圖來概括，就是下面這樣：

這就引入一個新的話題——Function類是如何抽象出來的。

Function基類定義

Function的數(shù)據(jù)成員如下所示：

using edge_list = std::vector;
using variable_list = std::vector;

struct TORCH_API Function {
...
  virtual variable_list apply(variable_list&& inputs) = 0;
...
  const uint64_t sequence_nr_;
  edge_list next_edges_;
  PyObject* pyobj_ = nullptr; // weak reference
  std::unique_ptr anomaly_metadata_ = nullptr;
  std::vector> pre_hooks_;
  std::vector> post_hooks_;
  at::SmallVector input_metadata_;
};

Function call

Function類是抽象出來的基類，代表一個op（operation），每個op接收的參數(shù)是0個、1個或多個Variable實例（使用std::vector封裝），并與此同時輸出0個、1個或多個Variable實例。PyTorch中所有用于反向傳播計算的函數(shù)都繼承自Function類，并重寫了Function類中的apply純虛函數(shù)。因為Function類中實現(xiàn)了call函數(shù)：

variable_list operator()(variable_list&& inputs) {
  return apply(std::move(inputs));
}

所以依靠C++的多態(tài)，對op的call將轉(zhuǎn)化為自身（子類）的apply調(diào)用。Function類中最重要的方法是call函數(shù)，call會調(diào)用apply，call函數(shù)接收vector封裝的多個Variable實例，并輸出vector封裝的多個Variable實例。輸入?yún)?shù)的vector長度可以由num_inputs()調(diào)用獲得，對應(yīng)的，輸出的vector長度則由num_outputs()獲得。

Function的輸入

Function成員input_metadata_代表input data的meta信息，界定了一個Function的輸入：

struct InputMetadata {
...
  const at::Type* type_ = nullptr;
  at::DimVector shape_;
  at::Device device_ = at::kCPU;
};

Autograd graph的edge和vertices

如果將PyTorch的autograd系統(tǒng)看作是一個圖（graph）的話，那么每個Function實例就是graph中的節(jié)點（nodes/vertices），各個Function實例之間則是通過Edge連接的。Edge是個結(jié)構(gòu)體，通過 (Function, input_nr) 的配對來代表graph中的edge：

struct Edge {
...
  std::shared_ptr function;
  uint32_t input_nr;
};

Function的成員next_edges_正是一組這樣的Edge實例，代表此function實例的返回值要輸出到的（另外）function，也即next_edges_是function和function之間的紐帶。

Function的輸入輸出都是Variable實例，因此，當一個graph被執(zhí)行的時候，Variable實例就在這些edges之間來傳輸流動。當兩個或者多個Edge指向同一個Function的時候（這個節(jié)點的入度大于1），這些edges的輸出將會隱含的相加起來再送給指向的目標Function。

Function和Function之間通過next_edge接口連接在一起，你可以使用add_next_edge()來向Function添加一個edge, 通過next_edge(index)獲取對應(yīng)的edge，通過next_edges()方法獲得迭代edge的迭代器。每一個Function都有一個sequence number，隨著Function實例的不斷構(gòu)建而單調(diào)增長。你可以通過sequence_nr()方法來或者一個Function的sequence number。

Function繼承體系

基類Function直接派生出TraceableFunction和以下這些Function：

CopySlices : public Function 
DelayedError : public Function 
Error : public Function 
Gather : public Function 
GraphRoot : public Function 
Scatter : public Function
AccumulateGrad : public Function 
AliasBackward : public Function 
AsStridedBackward : public Function 
CopyBackwards : public Function 
DiagonalBackward : public Function 
ExpandBackward : public Function 
IndicesBackward0 : public Function 
IndicesBackward1 : public Function 
PermuteBackward : public Function 
SelectBackward : public Function 
SliceBackward : public Function 
SqueezeBackward0 : public Function 
SqueezeBackward1 : public Function 
TBackward : public Function 
TransposeBackward0 : public Function 
UnbindBackward : public Function 
UnfoldBackward : public Function 
UnsqueezeBackward0 : public Function 
ValuesBackward0 : public Function 
ValuesBackward1 : public Function 
ViewBackward : public Function

PyFunction : public Function

這其中，從基類Function派生出來的AccumulateGrad、TraceableFunction、GraphRoot是比較關(guān)鍵的類。

派生類AccumulateGrad

先說說AccumulateGrad，AccumulateGrad正是Variable的grad_accumulator_成員的類型：

struct?AccumulateGrad?:?public?Function?{
??explicit?AccumulateGrad(Variable?variable_);
??variable_list?apply(variable_list&&?grads)?override;
??Variable?variable;
};

可見一個AccumulateGrad實例必須用一個Variable構(gòu)建，apply調(diào)用接收一個list的Variable的實例——這都是和Variable的grad_accumulator_相關(guān)的。

派生類GraphRoot

對于GraphRoot，前向時候的最終輸出——在反向的時候作為最初輸入——是由GraphRoot封裝的：

struct?GraphRoot?:?public?Function?{
??GraphRoot(edge_list?functions,?variable_list?inputs)
??????:?Function(std::move(functions)),
????????outputs(std::move(inputs))?{}
??variable_list?apply(variable_list&&?inputs)?override?{
????return?outputs;
??}
??variable_list?outputs;
};

GraphRoot——正如Function的靈魂在apply一樣——其apply函數(shù)僅僅返回它的輸入！

派生類TraceableFunction

再說說TraceableFunction：

struct TraceableFunction : public Function {
  using Function::Function;
  bool is_traceable() final {
    return true;
  }
};

TraceableFunction會進一步派生出372個子類（2019年4月），這些子類的名字都含有一個共同的部分：Backward。這說明什么呢？這些函數(shù)將只會用在反向傳播中：

AbsBackward : public TraceableFunction 
AcosBackward : public TraceableFunction 
AdaptiveAvgPool2DBackwardBackward : public TraceableFunction 
AdaptiveAvgPool2DBackward : public TraceableFunction 
AdaptiveAvgPool3DBackwardBackward : public TraceableFunction 
AdaptiveAvgPool3DBackward : public TraceableFunction 
AdaptiveMaxPool2DBackwardBackward : public TraceableFunction 
AdaptiveMaxPool2DBackward : public TraceableFunction 
AdaptiveMaxPool3DBackwardBackward : public TraceableFunction 
AdaptiveMaxPool3DBackward : public TraceableFunction 
AddBackward0 : public TraceableFunction 
AddBackward1 : public TraceableFunction 
AddbmmBackward : public TraceableFunction 
AddcdivBackward : public TraceableFunction 
AddcmulBackward : public TraceableFunction 
AddmmBackward : public TraceableFunction 
AddmvBackward : public TraceableFunction 
AddrBackward : public TraceableFunction 
......
SoftmaxBackwardDataBackward : public TraceableFunction 
SoftmaxBackward : public TraceableFunction 
......
UpsampleBicubic2DBackwardBackward : public TraceableFunction 
UpsampleBicubic2DBackward : public TraceableFunction 
UpsampleBilinear2DBackwardBackward : public TraceableFunction 
UpsampleBilinear2DBackward : public TraceableFunction 
UpsampleLinear1DBackwardBackward : public TraceableFunction 
UpsampleLinear1DBackward : public TraceableFunction 
UpsampleNearest1DBackwardBackward : public TraceableFunction 
UpsampleNearest1DBackward : public TraceableFunction 
UpsampleNearest2DBackwardBackward : public TraceableFunction 
UpsampleNearest2DBackward : public TraceableFunction 
UpsampleNearest3DBackwardBackward : public TraceableFunction 
UpsampleNearest3DBackward : public TraceableFunction 
UpsampleTrilinear3DBackwardBackward : public TraceableFunction 
UpsampleTrilinear3DBackward : public TraceableFunction 
......

這300多個Backward function都重寫了apply函數(shù)，來實現(xiàn)自己的反向求導(dǎo)算法，比如加法的反向求導(dǎo)函數(shù)AddBackward0：

struct AddBackward0 : public TraceableFunction {
  using TraceableFunction::TraceableFunction;
  variable_list apply(variable_list&& grads) override;
  Scalar alpha;
};

這些apply函數(shù)是Function的靈魂，是反向傳播計算時候的核心執(zhí)行邏輯。

Engine

Engine類實現(xiàn)了從輸出的variable（以及它的gradients）到root variables（用戶創(chuàng)建的并且requires_grad=True）之間的反向傳播。

gemfield = torch.ones(2, 2, requires_grad=True)
syszux = gemfield + 2
civilnet = syszux * syszux * 3
gemfieldout = civilnet.mean()
gemfieldout.backward()

還是以上面這個代碼片段為例，Engine實現(xiàn)了從gemfieldout到gemfield的反向傳播：

1，如何根據(jù)gemfieldout構(gòu)建GraphRoot；

2，如何根據(jù)這些Function實例及它們上的metadata構(gòu)建graph；

3，如何實現(xiàn)Queue來多線程完成反向計算的工作。

Engine類定義

Engine類的定義如下：

struct Engine {
  using ready_queue_type = std::deque, InputBuffer>>;
  using dependencies_type = std::unordered_map;
  virtual variable_list execute(const edge_list& roots,const variable_list& inputs,...const edge_list& outputs = {});
  void queue_callback(std::function callback);
protected:
  void compute_dependencies(Function* root, GraphTask& task);
  void evaluate_function(FunctionTask& task);
  void start_threads();
  virtual void thread_init(int device);
  virtual void thread_main(GraphTask *graph_task);
  std::vector> ready_queues;
};

核心就是execute函數(shù)，它接收一組Edge——(Function, input number) pairs ——來作為函數(shù)的輸入，然后通過next_edge不斷的找到指向的下一個Edge，最終完成整個Graph的計算。

派生類PythonEngine

然而我們實際使用的是Engine類的派生類：PythonEngine。PythonEngine子類重寫了父類的execute，只不過僅僅提供了把C++異常翻譯為Python異常的功能，核心工作還是由Engine基類來完成：

struct PythonEngine : public Engine

整個PyTorch程序全局只維護一個Engine實例，也就是PythonEngine實例。

BP調(diào)用棧

既然Engine是用來計算網(wǎng)絡(luò)反向傳播的，我們不妨看下這個調(diào)用棧是怎么到達Engine類的。如果我們對gemfieldout進行backward計算，則調(diào)用棧如下所示：

#torch/tensor.py，self is gemfieldout
def backward(self, gradient=None, retain_graph=None, create_graph=False)
|
V
#torch.autograd.backward(self, gradient, retain_graph, create_graph)
#torch/autograd/__init__.py
def backward(tensors, grad_tensors=None, retain_graph=None, create_graph=False, grad_variables=None)
|
V
Variable._execution_engine.run_backward(tensors, grad_tensors, retain_graph, create_graph,allow_unreachable=True)
#轉(zhuǎn)化為Variable._execution_engine.run_backward((gemfieldout,), (tensor(1.),), False, False,True)
|
V
#torch/csrc/autograd/python_engine.cpp
PyObject *THPEngine_run_backward(THPEngine *self, PyObject *args, PyObject *kwargs)
|
V
#torch/csrc/autograd/python_engine.cpp
variable_list PythonEngine::execute(const edge_list& roots, const variable_list& inputs, bool keep_graph, bool create_graph, const edge_list& outputs)
|
V
#torch/csrc/autograd/engine.cpp
Engine::execute(roots, inputs, keep_graph, create_graph, outputs)

總結(jié)

在下段文章中，Gemfield將主要介紹Engine這個類是如何在gemfieldout.backward()中運行PyTorch動態(tài)圖的。

PyTorch的動態(tài)圖（下）

背景

在上文中，我們介紹了PyTorch autograd系統(tǒng)的三個基石：Variable、Function、Engine。

用一句簡單的話來概括下，就是Engine使用Function構(gòu)建的Graph來計算得到Variable上grad。在本文中，Gemfield將以下面的代碼片段為例，詳細介紹Engine如何構(gòu)建Graph來進行反向傳播計算：

gemfield = torch.ones(2, 2, requires_grad=True)
syszux = gemfield + 2
civilnet = syszux * syszux * 3
gemfieldout = civilnet.mean()
gemfieldout.backward()

BP Engine

BP Engine是一個反向傳播計算中用于動態(tài)生成計算圖的類，目前PyTorch中只定義了一種BP Engine的實現(xiàn)。

1，Engine類定義

用于反向傳播計算圖動態(tài)生成的Engine類的定義如下：

struct Engine {
  using ready_queue_type = std::deque, InputBuffer>>;
  using dependencies_type = std::unordered_map;
  virtual variable_list execute(const edge_list& roots,const variable_list& inputs,...const edge_list& outputs = {});
  void queue_callback(std::function callback);
protected:
  void compute_dependencies(Function* root, GraphTask& task);
  void evaluate_function(FunctionTask& task);
  void start_threads();
  virtual void thread_init(int device);
  virtual void thread_main(GraphTask *graph_task);
  std::vector> ready_queues;
};

前向結(jié)束輸出gemfieldout后，我們使用gemfieldout來作為反向傳播的輸入。

2，Engine類的start_threads成員

顧名思義，start_threads就是用來啟動線程的，根據(jù)設(shè)備的數(shù)量來決定要啟動的線程數(shù)量。這個函數(shù)使用std::call_once(start_threads_flag, &Engine::start_threads, this)的方式進行調(diào)用，確保了整個進程周期start_threads成員函數(shù)只被調(diào)用了一次。該成員函數(shù)的主要作用有2點：

1，創(chuàng)建多個ReadyQueue實例，使用ready_queues vector來管理。ReadyQueue的數(shù)量和即將要新創(chuàng)建的線程數(shù)量一樣：

ready_queues = std::vector>(num_threads);
for (auto& queue : ready_queues){
    queue.reset(new ReadyQueue());
}

2，創(chuàng)建多個新線程，線程的數(shù)量取決于設(shè)備的數(shù)量。CPU算一個設(shè)備，每張GPU卡算一個設(shè)備，然后最后再加1。比如該系統(tǒng)上有4個RTX 2080ti顯卡，那么這里就會啟動5個線程，如果系統(tǒng)上只有cpu而沒有GPU，那么這里就會啟動2個線程。

該成員函數(shù)使用std::thread來啟動管理線程，比較重要的一點是，創(chuàng)建線程的時候傳遞的this指針：

for?(int?i?=?0;?i?????std::thread?t(&Engine::thread_init,?this,?i?-?1);
????t.detach();
}

this就是當前Engine的實例，整個進程的生命周期內(nèi)就只有這一個Engine實例。傳遞this帶來的一個驚喜就是，當前進程和新啟動的線程之間可以共享同一個Engine實例——不管是數(shù)據(jù)成員還是函數(shù)成員。后面你就會看到，我們的Queue就依靠this的共享來實現(xiàn)線程間的對象傳輸。

3，Engine類的ready_queues

ready_queues的定義如下：

std::vector> ready_queues;

可見ready_queues使用vector去管理了若干個ReadyQueue實例。這樣我們就可以使用device index去vector里索引每個device專屬的ReadyQueue了。

ReadyQueue用來傳輸?shù)氖荈unctionTask對象（后文會介紹到），ReadyQueue定義如下所示：

struct ReadyQueue {
  std::priority_queue, CompareFunctionTaskTime> heap;
  std::condition_variable not_empty;
  std::mutex mutex;
  void push(FunctionTask item);
  FunctionTask pop();
};

ReadyQueue使用priority_queue作為backend這一事實告訴了我們一點：消費的順序并不等價于生產(chǎn)的順序——根據(jù)CompareFunctionTaskTime的定義——誰的sequence_nr()越小就誰先消費。

ReadyQueue使用了C++11的condition_variable進行線程間的同步，使用condition_variable的notify_one來通知消費線程，相當于unblock其中一個消費線程（notify_all則是所有消費線程）。與此對應(yīng)，消費線程則使用condition_variable的wait來接收同步信息。ReadyQueue類上定義了push和pop方法，分別代表生產(chǎn)者的生產(chǎn)行為和消費者的消費行為：

auto ReadyQueue::push(FunctionTask item) -> void {
  {
    std::lock_guard lock(mutex);
    ++item.base->outstanding_tasks;
    heap.push(std::move(item));
  }
  not_empty.notify_one();
}

auto ReadyQueue::pop() -> FunctionTask {
  std::unique_lock lock(mutex);
  not_empty.wait(lock, [this]{ return !heap.empty(); });
  auto task = std::move(const_cast(heap.top()));
  heap.pop();
  return task;
}

//wait相當于下面，防止異常情況退出
while(heap.empty()){
    not_empty.wait(lock);
}

4，Engine類的thread_init成員

thread_init會執(zhí)行start_threads啟動的線程的初始化工作。

auto Engine::thread_init(int device) -> void {
  at::init_num_threads();
  std::array(c10::DeviceType::COMPILE_TIME_MAX_DEVICE_TYPES)> guards;
  if (device != -1) {
    for (size_t i = 0; i < static_cast(c10::DeviceType::COMPILE_TIME_MAX_DEVICE_TYPES); i++) {
      auto* impl = c10::impl::device_guard_impl_registry[i].load();
      if (impl && device < impl->deviceCount()) {
        guards[i].reset_device(at::Device(static_cast(i), device));
      }
    }
  }
  worker_device = device;
  thread_main(nullptr);
}

設(shè)置每個線程自己的worker_device的值，值為device num-1，因此是從-1開始的。如果系統(tǒng)上只有1個cpu設(shè)備，那么start_threads會啟動2個線程。那么這里的device號就分別是-1、0。另外，主進程里的worker_device的值為NO_DEVICE（-2）。初始化工作除了設(shè)置這個worker_device，主要是從第2個worker thread開始來設(shè)置guards中的divice。線程初始化完畢后，將會進行真正的線程執(zhí)行邏輯調(diào)用。

5，Engine類中的GraphTask

Engine中使用的GraphTask：

struct GraphTask {
  std::atomic outstanding_tasks;
  bool keep_graph;
  bool grad_mode;
  std::mutex mutex;
  std::condition_variable not_done;
  std::unordered_map not_ready;
  std::unordered_map dependencies;
  struct ExecInfo {
    bool needed = false;
  };
  std::unordered_map exec_info;
  int owner;
  GraphTask(bool keep_graph, bool grad_mode): has_error(false),  \
            outstanding_tasks(0), keep_graph(keep_graph), grad_mode(grad_mode), owner(NO_DEVICE) {}
};

在Engine的execute函數(shù)執(zhí)行中，我們會定義一個graph_task實例：

GraphTask?graph_task(keep_graph,?create_graph);

GraphTask的重要成員有：

成員1：outstanding_tasks，是個數(shù)字，當一個GraphTask實例創(chuàng)建出來的時候，outstanding_tasks被初始化為0；當其隨后被送入ReadyQueue的時候，outstanding_tasks自增1；然后在worker線程每執(zhí)行一次evaluate_function(task)后，outstanding_tasks的值減1。并且在主進程中，會有一個線程的同步邏輯依賴這個值：

while(graph_task.outstanding_tasks.load() != 0){
    graph_task.not_done.wait(lock);
}

可見這個graph_task實例上的function沒有被evaluate完成之前，主進程就一直在這里等待。

成員2：keep_graph，是個bool值，用來指定一次反向計算后是否釋放資源。什么資源呢？前向過程中建立起來的資源。keep_graph如果是False的話，則會在fn執(zhí)行完畢后調(diào)用release_variables：

if (!task.base->keep_graph) {
  fn.release_variables();
}

我們在上文提到過，那幾百個反向計算的Function里都有一個靈魂函數(shù)——apply，其實，還有一個回收資源的靈魂函數(shù)——release_variables。比如

struct MulBackward0 : public TraceableFunction {
  void release_variables() override {
    self_.reset_data();
    self_.reset_grad_function();
    other_.reset_data();
    other_.reset_grad_function();
  }
  SavedVariable self_;
  SavedVariable other_;
};

成員3：grad_mode，這是個bool值，用來指示當前的上下文是否是要計算grad。

bool GradMode::is_enabled() {
  return GradMode_enabled;
}

void GradMode::set_enabled(bool enabled) {
  GradMode_enabled = enabled;
}

整個反向計算期間執(zhí)行的代碼邏輯中，都是靠GradMode::is_enabled()來判斷當前是否是要計算grad的。

成員4：mutex，類型為std::mutex。這是線程之間的同步原語，多個線程中只有一個可以持有同一個mutex，其它的線程只能在此等待。但是為了防止死鎖等情況出現(xiàn)，我們借助RAII來智能管理mutex，典型的代表就是std::lock_guard和std::unique_lock。默認情況下用std::lock_guard，在一個函數(shù)執(zhí)行完畢后（其實是一個block作用域），std::lock_guard的析構(gòu)會自動釋放這個mutex；不過，std::lock_guard只能靠構(gòu)造函數(shù)和析構(gòu)函數(shù)來獲得和釋放mutex，std::unique_lock則在此基礎(chǔ)上更近一步，除了擁有上述功能之外，還可以通過std::unique_lock的lock和unlock來主動加鎖和解鎖，提供更精細粒度的控制——增加代碼并行的區(qū)域。

成員5：not_done，std::condition_variable類型。線程間通信。還記得吧，在主進程中，我們將一個FunctionTask對象送往Queue后，主進程就開始等待：

while(graph_task.outstanding_tasks.load() != 0){
    graph_task.not_done.wait(lock);
}

while循環(huán)就是防止異常退出的，核心是not_done這個condition。not_done.wait(lock)就是在這里阻塞——等待worker thread里的not_done.notify_all()。

成員6：not_ready，類型為std::unordered_map。用來暫存not ready的function及其輸入。

**成員7：dependencies，**類型為std::unordered_map 。

這個實例的dependencies成員在compute_dependencies調(diào)用中被初始化，只要一個grad_fn函數(shù)在別人的next_edges()中出現(xiàn)過一次，那么dependencies[this_grad_fn] 就自增1。

成員8：exec_info，類型為std::unordered_map 。如果exec_info這個map為空的話，說明這個task是默認的模式——所有我們在next_edges中遇到的函數(shù)都將被執(zhí)行。如果exec_info不為空的話，只有含有entry的Function并且這個entry
has needed == True 的情況下才會被執(zhí)行。

成員9：owner，int類型。GraphTask是在哪個線程中創(chuàng)建的，該值就是那個線程中的worker_device的值。

6，Engine類中的FunctionTask

Engine中使用的FunctionTask：

struct FunctionTask {
  GraphTask* base;
  std::shared_ptr fn;
  InputBuffer inputs;
  FunctionTask(GraphTask* base, std::shared_ptr fn, InputBuffer inputs): \
               base(base), fn(std::move(fn)), inputs(std::move(inputs)) {}
};

這個類的對象正是在queue中傳輸?shù)臇|西，從上面的定義可以看到，我們使用GraphTask、Function、InputBuffer來構(gòu)建一個FunctionTask實例。實際上，在主進程中，我們只需要將一個FunctionTask實例送入queue即可。這個FunctionTask正是由上述的GraphTask實例、graph_root、0構(gòu)建的：

#主進程中
FunctionTask(&graph_task, std::move(graph_root), InputBuffer(0)

graph_root初始化很簡單，由roots和inputs構(gòu)建，roots就是將gemfieldout的gradient_edge()——也即grad_fn——也即MeanBackward0實例和output_nr_——也即（MeanBackward0實例，0）；而inputs也即tensor(1.)。生產(chǎn)端往queue發(fā)送了FunctionTask實例后，在消費端的worker thread中，我們則通過task.base來訪問到這個GraphTask實例、通過task.fn訪問到這個roots實例、通過task.inputs來訪問這個InputBuffer實例。

而在worker thread中，我們也可能使用如下的方式來構(gòu)建新的FunctionTask實例，然后添加到queue中：

#work thread
FunctionTask(task.base, nullptr, InputBuffer(0)

#evaluate function
FunctionTask(task.base, next.function, std::move(input_buffer))

7，Engine類中的compute_dependencies

還記得5中定義的那個graph_task實例嗎？還記得graph_task中有個成員是dependencies嗎？還記得dependencies是std::unordered_map嗎？對，只要一個grad_fn函數(shù)在別人的next_edges()中出現(xiàn)過一次，那么dependencies[this_grad_fn] 就自增1。

對，這個函數(shù)就只干了這么一件事。

8，Engine類中的execute

這是Engine的靈魂函數(shù)，也是Engine的執(zhí)行邏輯的主體，在一次反向中，該函數(shù)會被執(zhí)行一次。該函數(shù)主要做了如下的工作：

1，調(diào)用Engine::start_threads，啟動多個worker線程；注意，start_threads在整個進程周期只會被執(zhí)行一次；

2，實例化一個graph_task的實例；

3，使用上述graph_task的mutex，初始化一個execute函數(shù)中局部的mutex；

4，構(gòu)建GraphRoot；

5，執(zhí)行compute_dependencies；

6，向隊列中傳輸一個FunctionTask實例；

7，等待worker thread計算完畢后結(jié)束。

9，Engine類的thread_main成員

thread_main會作為start_threads啟動的線程的執(zhí)行實體：

auto Engine::thread_main(GraphTask *graph_task) -> void {
  auto queue = ready_queues[worker_device + 1];
  while (!graph_task || graph_task->outstanding_tasks > 0) {
    FunctionTask task = queue->pop();
    if (task.fn && !task.base->has_error.load()) {
      GradMode::set_enabled(task.base->grad_mode);
      evaluate_function(task);
    }
    auto base_owner = task.base->owner;

    if (base_owner == NO_DEVICE) {
      if (--task.base->outstanding_tasks == 0) {
        std::lock_guard lock(task.base->mutex);
        task.base->not_done.notify_all();
      }
    } else {
      if (base_owner == worker_device) {
        --task.base->outstanding_tasks;
      } else if (base_owner != worker_device) {
        if (--task.base->outstanding_tasks == 0) {
          std::atomic_thread_fence(std::memory_order_release);
          ready_queue_by_index(base_owner).push(FunctionTask(task.base, nullptr, InputBuffer(0)));
        }
      }
    }
  }
}

主體就是一個while循環(huán)，不斷的從專屬queue中取出FunctionTask實例，然后執(zhí)行evaluate_function——這是主體邏輯。evaluate_function完畢后，如果：

1，這個task是來自主進程的，并且task的outstanding_tasks已經(jīng)降為0了（注意outstanding_tasks在evaluate_function中還會不斷的被改變），那么就通知主進程上的wait同步原語，準備結(jié)束反向計算；

2，如果這個task是來自當前work thread，那么outstanding_tasks就自減1；

3，如果這個task是來自其它work thread，那么outstanding_tasks就自減1，并且如果降為0的話，則向那個worker thread的queue發(fā)送一個dummy function task。

10，Engine類的evaluate_function

這個函數(shù)的核心就是看一個函數(shù)在GraphTask的dependencies的計數(shù)是否降為0——也即是否ready（通常是因為有多個input），如果ready就送往queue去計算；如果沒有就放入GraphTask的not_ready中，每次設(shè)置好對應(yīng)的一個InputBuffer輸入。送往Queue的task和之前GraphRoot不一樣，因為它的owner不是主進程——是在worker thread中創(chuàng)建的。

1，一個node的input variable的device是多少，那么這個node組成的FunctionTask對象將被送往那個device對應(yīng)的worker thread的queue；

2，準備一個node的時候，如果is_ready是True，說明這個node不會被未來的計算所依賴了，這個node組成的FunctionTask就會被送往上述的queue，并且信息會從Graph（GraphTask）中抹掉——從GraphTask的dependencies中移除；

3，準備一個node的時候，如果is_ready是False，這通常表明這個node有多個輸入（被更多的node連接，使用num_inputs()可以獲得數(shù)量），那么第一次遇到這個node的時候并不會把它送往queue，而是放入GraphTask的not_ready中，這個時候順便設(shè)置好了該node的第一個輸入：

input_buffer.add(next.input_nr, std::move(output))
not_ready.emplace(next.function.get(), std::move(input_buffer));

第二次遇到該node就設(shè)置該node的第二個input：

auto &input_buffer = not_ready_it->second;
input_buffer.add(next.input_nr, std::move(output));

以此類推，直到dependencies中對應(yīng)的計數(shù)降為0——is_ready從False變?yōu)榱薚rue——這個node終于不會被未來的計算所依賴了。這個node組成的FunctionTask這個時候才會被送往上述的queue，并且信息會從Graph（GraphTask）中抹掉——從GraphTask的dependencies中移除、從GraphTask的not_ready中移除。

另外，input_buffer的構(gòu)建也非常有趣，當向一個node添加一個輸入的時候：

input_buffer.add(next.input_nr, std::move(output));

input_nr就清楚的表明，當前的node是反向傳播中要流向的node的第幾個輸入！

11，Engine類的call_function

這個函數(shù)的邏輯就相對簡單了，調(diào)用各個反向計算的函數(shù)及注冊在上面的hooks。注意fn的輸入和輸出，像gemfield在《PyTorch的Tensor》中說的那樣，輸入是一組Variable的實例——InputBuffer::variables(std::move(task.inputs))，輸出也是一組Variable的實例，畢竟輸出要作為下一個fn的輸入嘛。相關(guān)的代碼如下所示：

static variable_list call_function(FunctionTask& task) {
  auto& fn = *task.fn;
  auto inputs = call_pre_hooks(fn, InputBuffer::variables(std::move(task.inputs)));

  const auto has_post_hooks = !fn.post_hooks().empty();
  variable_list outputs = fn(std::move(inputs));
  
  if(has_post_hooks){
    return call_post_hooks(fn, std::move(outputs), inputs);
  }
  return outputs;
}

1，調(diào)用注冊在node上的pre_hooks；

2，調(diào)用node本身，比如MeanBackward0、MulBackward0等；

3，調(diào)用注冊在node上的post hooks。

總結(jié)

本文中g(shù)emfield進一步介紹了PyTorch的動態(tài)圖，主要是Engine這個class。現(xiàn)在你已經(jīng)熟悉了，反向傳播計算中，會根據(jù)設(shè)備數(shù)量啟動多個線程，每個線程關(guān)聯(lián)一個Queue，每個worker thread都有關(guān)聯(lián)的device，主進程和worker thread、worker thread和work thread之間通過queue傳送輸入和計算結(jié)果，task根據(jù)input的device會被送入對應(yīng)的Queue——如此以來計算就會在那個device關(guān)聯(lián)的worker thread中進行。數(shù)據(jù)流經(jīng)每個fn/node獲得的輸出再作為下一個fn/node的輸入，這些輸入輸出都是一組Variable的實例。

推薦閱讀

Facebook AI PyTorch數(shù)據(jù)并行訓練秘籍大揭秘
總結(jié)深度學習PyTorch神經(jīng)網(wǎng)絡(luò)箱使用
Pytorch轉(zhuǎn)Msnhnet模型思路分享

?ACCV 2020國際細粒度網(wǎng)絡(luò)圖像識別競賽正式開賽！

添加極市小助手微信（ID : cvmart2），備注：姓名-學校/公司-研究方向-城市（如：小極-北大-目標檢測-深圳），即可申請加入極市目標檢測/圖像分割/工業(yè)檢測/人臉/醫(yī)學影像/3D/SLAM/自動駕駛/超分辨率/姿態(tài)估計/ReID/GAN/圖像增強/OCR/視頻理解等技術(shù)交流群：每月大咖直播分享、真實項目需求對接、求職內(nèi)推、算法競賽、干貨資訊匯總、與?10000+來自港科大、北大、清華、中科院、CMU、騰訊、百度等名校名企視覺開發(fā)者互動交流~

△長按添加極市小助手

△長按關(guān)注極市平臺，獲取最新CV干貨

覺得有用麻煩給個在看啦~??

可能是最詳盡的PyTorch動態(tài)圖解析

PyTorch的動態(tài)圖（上）

背景

autograd初始化

Variable

Function

Function繼承體系

Engine

BP調(diào)用棧

總結(jié)

背景

BP Engine

總結(jié)