[源碼解析] PyTorch如何實現前向傳播(3) --- 具體實現
0x00 摘要
本系列將通過大概十篇左右文章來分析 PyTorch 的自動微分功能如何實現。本文是前向傳播的第三篇，介紹具體實現機制。
在反向傳播時候，當拿到了一個張量，引擎需要知道：
- 如何對此張量調用梯度計算，即從哪里找到計算梯度的函數 F。
- 拿到函數 F 之后，這個函數的輸入就是此張量本身，但是函數 F 需要知道輸入參數（本張量）的一些元信息，比如類型，shape，device。
- F 計算出梯度之后，需要知道 F 的輸出應該傳播到哪里，就是怎么在反向傳播計算圖上繼續進行下一步。
本文就是具體分析，在前向傳播之中這些信息如何設置。
本系列前幾篇連接如下：
[源碼解析]PyTorch如何實現前向傳播(1) --- 基礎類(上)
[源碼解析]PyTorch如何實現前向傳播(2) --- 基礎類(下)
0x01 計算圖
1.1 圖的相關類
計算圖是一個有向圖，它的節點為已經實現的算子或者數據（葉子結點），箭頭的方向表示數據流動的方向，從輸入節點指向輸出節點。由前面章節可知，圖相關有三個基本類：Node，Edge，Engine（我們后續會分析Engine）。
- 節點是 Node 類，代表一個操作（operation）。
- 每個 Node 接收0個或者多個 Variable，輸出0個或者多個 Variable。Node 之間由 Edge 連接在一起，其實就是通過 Node 的成員變量
next_edges_連接在一起。 - 反向傳播的函數都繼承自 Node，比如 SubBackward0就繼承自 Node。
- 每個 Node 接收0個或者多個 Variable，輸出0個或者多個 Variable。Node 之間由 Edge 連接在一起，其實就是通過 Node 的成員變量
- 邊 Edge 其實本質是 (Node, input_nr)
- Edge 的成員變量 std::shared_ptr
function ：指定本邊指向的Node。 - Edge 的成員變量 uint32_t input_nr ： 指定本邊是function的第幾個輸入 。
- Edge 的成員變量 std::shared_ptr
- Node 的成員 next_edges_ 是一組 Edge實例，代表此 Node 實例的返回值要輸出到的（另外）Node，即 next_edges_是 Node 和Node 之間的紐帶。當計算圖被執行時候，Variable 在這些邊之間流動。
- Engine 是執行引擎。
1.2 動態圖
pytorch在設計中采取了動態計算圖的方式。動態的意思是：反向傳播的計算圖是動態更新的。每一輪反向傳播開始時（前向傳播結束后）都會動態的重新構建一個計算圖，當本次反向傳播完成后，圖會銷毀計算圖，在內存中被釋放了。如果在新一輪中想再次使用，只能從頭再搭建一遍。這種動態更新的方式允許用戶在迭代過程中更改網絡的形狀和大小。
下面代碼可以看出來動態圖的特質。
# 第一遍，生成動態圖
a = torch.tensor(2., requires_grad=True)
b = torch.tensor(6., requires_grad=True)
Q = 3*a**3 - b**2
external_grad = torch.tensor(1.)
Q.backward(gradient=external_grad) # 正常
Q.backward(gradient=external_grad) # RuntimeError
# 第二次：再來一遍
a = torch.tensor(2., requires_grad=True)
b = torch.tensor(6., requires_grad=True)
Q = 3*a**3 - b**2
external_grad = torch.tensor(1.)
Q.backward(gradient=external_grad) # 正常
1.3 動態展示
下面是PyTorch 官方的動態圖，大家可以有一個形象的理解。
為了更好的展示，我們把動圖分解開來看。
首先是聲明了一些張量。
其次讓兩個矩陣相乘。
讓另外兩個矩陣相乘
然后把兩個相乘結果相加。
加入 Tanh 激活函數。
加入損失函數。
反向傳播，計算梯度。
由此可以看出來，動態圖關系是在前向計算過程中構建出來的。
0x02 總體分析
我們把前文提到的示例代碼繼續細化，目的是為了看看計算圖中各個張量：
a = torch.tensor(2., requires_grad=True)
b = torch.tensor(6., requires_grad=True)
X = a ** 3
Y = 3 * X
Z = b ** 2
Q = X - Z
external_grad = torch.tensor(1.)
Q.backward(gradient=external_grad)
print(a.grad)
看看運行時變量如下，因為 Q = X - Z 是減法，所以對應的反向操作就是 SubBackward0：
Q = {Tensor} tensor(-28., grad_fn=<SubBackward0>)
X = {Tensor} tensor(8., grad_fn=<PowBackward0>)
Y = {Tensor} tensor(24., grad_fn=<MulBackward0>)
Z = {Tensor} tensor(36., grad_fn=<PowBackward0>)
a = {Tensor} tensor(2., requires_grad=True)
b = {Tensor} tensor(6., requires_grad=True)
我們可以和DAG 的可視化表示比對一下。在圖中，箭頭指向前向傳遞的方向，節點代表前向傳遞中每個操作的后向函數。藍色的葉子節點 (2) 代表我們的葉子張量a和b。
在代碼層面，在正向傳播過程中，PyTorch 並沒有顯式構造出一個反向傳播的計算圖，而是建立了若干所需的數據結構，可以認為是一個虛擬圖關系，但是沒有真實的圖數據結構。在每次迭代的前向傳播中，針對 Q = X - Z，都會執行如下操作：
- 1）進入減法操作 ：減法操作會派發到某一個device之上，其中會進行 Q 的構建。
- 2）先構建如何反向傳播 ：派發到 VariableType上時，會先進行 Q 的 autograd 信息的構建；
- 構建一個減法的反向計算函數 SubBackward0 實例。
- 初始化SubBackward0實例的
next_edges_和其它相關成員，next_edges_成員的值來自前向傳播的輸入參數 X 和 Z。- 如果輸入
Variable是leaf節點，則next_edges_來自輸入Variable的grad_accumulator_ - 如果輸入 Varaible是 非leaf節點，則
next_edges_來自輸入Variable的grad_fn_。
- 如果輸入
- 使用步驟 3 中的新Variable實例（就是前向計算的結果 Q）來初始化 SubBackward0 實例的
input_metadata_， - 這樣，就得到了如何進行 Q 的反向傳播，但此時只是得到了如何計算，還沒有和 Q 聯系起來。
- 3）再將 前向計算 & 與反向傳播 聯系起來 ：前向運算之后得到新的Variable，這個就是 Q，使用步驟2) 中的 SubBackward0 實例初始化 Q 的
autograd_meta_->grad_fn_成員。當對 Q 反向計算時候，就知道使用 Q 的autograd_meta_->grad_fn_成員來進行，就是 2) 之中的 SubBackward0。
大致如下圖所示：
+-----------------------+ +---------------------------------------+
| Q | | DifferentiableViewMeta |
| | | |
| autograd_meta_ +---------> | grad_ grad_accumulator_ |
| | | |
+-----------------------+ | |
+----------------------+ grad_fn_ output_nr_ | Q 找到如何計算梯度
| | |
| +---------------------------------------+
v
+-------------+------------+ +----------------------+
|SubBackward0 | | |
| | | Compute the gradient | 如何計算梯度
| apply +---------------> | |
| | +----------------------+
| |
| | +-----------------------------------------------------+
| next_edges_ +---------> | edge_list |
| | | |
| other_scalar_type | | [(PowBackward0(self), 0), (PowBackward0(other), 0)] | 輸出
| | | |
| alpha | +-----------------------------------------------------+
| |
| self_scalar_type | +----------------------------------------+
| | | |
| input_metadata_ +-----> | [(type of Q, shape of Q, device of Q)] | 輸入
| | | |
+--------------------------+ +----------------------------------------+
因為前向計算中會生成計算圖中的一系例節點，所以我們接下來就先分析這些節點。
0x03 Node 繼承體系
我們先從上圖中最下面的節點 SubBackward0 開始分析。
3.1 繼承體系
SubBackward0 定義位於：torch/include/torch/csrc/autograd/generated/Functions.h。
struct TORCH_API SubBackward0 : public TraceableFunction {
using TraceableFunction::TraceableFunction;
variable_list apply(variable_list&& grads) override;
std::string name() const override { return "SubBackward0"; }
void release_variables() override {
}
at::ScalarType other_scalar_type;
at::Scalar alpha;
at::ScalarType self_scalar_type;
};
我們再看看 SubBackward0 的繼承體系。
class SubBackward0 : public TraceableFunction
class TraceableFunction : public Node
/// See Node::is_traceable() for definition.
struct TraceableFunction : public Node {
using Node::Node;
bool is_traceable() final {
return true;
}
};
因此，SubBackward0 就是一個 Node 類型。
3.2 Node
前文我們已經介紹了 Node。Node 類，代表一個操作（operation）。每個 Node 接收0個或者多個 Variable，輸出0個或者多個 Variable。Node 之間由 Edge 連接在一起，其實就是通過 Node 的成員變量next_edges_連接在一起。反向傳播的函數都繼承自 Node。
我們提取部分 Node 代碼如下：
struct TORCH_API Node : std::enable_shared_from_this<Node> {
/// Performs the `Node`'s actual operation.
virtual variable_list apply(variable_list&& inputs) = 0;
const uint64_t sequence_nr_;
uint64_t topological_nr_ = 0;
// 在前向過程中與該算子相關聯的邊，對應了前向過程中的輸入variable。
edge_list next_edges_;
std::vector<std::unique_ptr<FunctionPreHook>> pre_hooks_;
std::vector<std::unique_ptr<FunctionPostHook>> post_hooks_;
at::SmallVector<InputMetadata, 2> input_metadata_;
// 這里對運算符()進行重載，核心其實就是調用apply()
variable_list operator()(variable_list&& inputs) {
bool pre_sampled = false;
if (at::shouldRunRecordFunction(&pre_sampled)) {
return apply(std::move(inputs));
} else {
return apply(std::move(inputs));
}
}
};
可以看到，apply(variable_list&& inputs) 是純虛函數，需要其派生類實現。 apply函數是Function的靈魂，是反向傳播計算時候的核心執行邏輯，通過 C++ 的多態功能就可以調用到各個派生類的 apply 函數。
3.3 SubBackward0
SubBackward0 的 apply函數代碼如下，可以看到其求導過程。代碼位於 torch/csrc/autograd/generated/Functions.cpp。
variable_list SubBackward0::apply(variable_list&& grads) {
IndexRangeGenerator gen;
auto self_ix = gen.range(1);
auto other_ix = gen.range(1);
variable_list grad_inputs(gen.size());
auto& grad = grads[0];
bool any_grad_defined = any_variable_defined(grads);
if (should_compute_output({ other_ix })) {
// 進行計算
auto grad_result = any_grad_defined ? (handle_r_to_c(other_scalar_type, -grad * alpha.conj())) : Tensor();
copy_range(grad_inputs, other_ix, grad_result); // 拷貝結果到grad_inputs
}
if (should_compute_output({ self_ix })) {
// 進行計算
auto grad_result = any_grad_defined ? (handle_r_to_c(self_scalar_type, grad)) : Tensor();
copy_range(grad_inputs, self_ix, grad_result); // 拷貝結果到grad_inputs
}
return grad_inputs; // 返回grad_inputs
}
我們印證一下，看看tools/autograd/derivatives.yaml 文件。這里是forward和backward的映射，可以理解為 autograd engine 在做反向鏈式求導時候查詢的原子操作，我們依據如下因此可以知道，加法和減法的求導函數都利用了 handle_r_to_c。
- name: add.Tensor(Tensor self, Tensor other, *, Scalar alpha=1) -> Tensor
self: handle_r_to_c(self.scalar_type(), grad)
other: handle_r_to_c(other.scalar_type(), maybe_multiply(grad, alpha.conj()))
result: self_t + maybe_multiply(other_t, alpha)
- name: sub.Tensor(Tensor self, Tensor other, *, Scalar alpha=1) -> Tensor
self: handle_r_to_c(self.scalar_type(), grad)
other: handle_r_to_c(other.scalar_type(), -grad * alpha.conj())
handle_r_to_c 定義如下，就是進行轉換。
Tensor handle_r_to_c(ScalarType self_st, Tensor gradient_result) {
if (!at::isComplexType(self_st) && gradient_result.is_complex()) {
// R -> C
return at::real(gradient_result);
}
return gradient_result;
}
用代碼來印證一下：
a = torch.tensor(2., requires_grad=True)
b = torch.tensor(6., requires_grad=True)
Q = a - b
external_grad = torch.tensor(1.)
Q.backward(gradient=external_grad)
這時候運行時如下：
a = {Tensor} tensor(2., requires_grad=True)
T = {Tensor} tensor(2., grad_fn=<PermuteBackward>)
data = {Tensor} tensor(2.)
device = {device} cpu
dtype = {dtype} torch.float32
grad = {Tensor} tensor(1.)
grad_fn = {NoneType} None
b = {Tensor} tensor(6., requires_grad=True)
T = {Tensor} tensor(6., grad_fn=<PermuteBackward>)
data = {Tensor} tensor(6.)
device = {device} cpu
dtype = {dtype} torch.float32
grad = {Tensor} tensor(-1.)
grad_fn = {NoneType} None
Q = {Tensor} tensor(-4., grad_fn=<SubBackward0>)
T = {Tensor} tensor(-4., grad_fn=<PermuteBackward>)
data = {Tensor} tensor(-4.)
device = {device} cpu
dtype = {dtype} torch.float32
grad = {NoneType} None
grad_fn = {SubBackward0} <SubBackward0 object at 0x7fb76e365438>
metadata = {dict: 0} {}
next_functions = {tuple: 2}
0 = {tuple: 2} (<AccumulateGrad object at 0x7fb76e344978>, 0)
1 = {tuple: 2} (<AccumulateGrad object at 0x7fb76e3447b8>, 0)
__len__ = {int} 2
requires_grad = {bool} True
is_cuda = {bool} False
is_leaf = {bool} False
is_meta = {bool} False
is_mkldnn = {bool} False
is_mlc = {bool} False
is_quantized = {bool} False
is_sparse = {bool} False
is_sparse_csr = {bool} False
is_vulkan = {bool} False
is_xpu = {bool} False
layout = {layout} torch.strided
name = {NoneType} None
names = {tuple: 0} ()
ndim = {int} 0
output_nr = {int} 0
requires_grad = {bool} True
shape = {Size: 0} torch.Size([])
我們接着看看其他幾個節點。
3.4 PowBackward0
PowBackward0 定義如下。
struct TORCH_API PowBackward0 : public TraceableFunction {
using TraceableFunction::TraceableFunction;
variable_list apply(variable_list&& grads) override;
std::string name() const override { return "PowBackward0"; }
void release_variables() override {
std::lock_guard<std::mutex> lock(mutex_);
self_.reset_data();
}
SavedVariable self_;
at::Scalar exponent;
};
variable_list PowBackward0::apply(variable_list&& grads) {
std::lock_guard<std::mutex> lock(mutex_);
IndexRangeGenerator gen;
auto self_ix = gen.range(1);
variable_list grad_inputs(gen.size());
auto& grad = grads[0];
auto self = self_.unpack();
bool any_grad_defined = any_variable_defined(grads);
if (should_compute_output({ self_ix })) {
auto grad_result = any_grad_defined ? (pow_backward(grad, self, exponent)) : Tensor();
copy_range(grad_inputs, self_ix, grad_result);
}
return grad_inputs;
}
我們去 tools/autograd/derivatives.yaml 中看看，看到就是使用了 pow_backward。
- name: pow.Tensor_Scalar(Tensor self, Scalar exponent) -> Tensor
self: pow_backward(grad, self, exponent)
result: auto_element_wise
最終也用到了handle_r_to_c。
Tensor pow_backward(Tensor grad, const Tensor & self, const Scalar & exponent) {
if (exponent.equal(0.0)) {
return at::zeros_like(self, LEGACY_CONTIGUOUS_MEMORY_FORMAT);
} else {
auto grad_lambda = [&](auto exp) { return grad * (exp * self.pow(exp - 1)).conj(); };
Tensor out = (exponent.isComplex()) ? grad_lambda(exponent.toComplexDouble()) : grad_lambda(exponent.toDouble());
return handle_r_to_c(self, out);
}
}
3.5 MulBackward0
MulBackward0 定義如下。
struct TORCH_API MulBackward0 : public TraceableFunction {
using TraceableFunction::TraceableFunction;
variable_list apply(variable_list&& grads) override;
std::string name() const override { return "MulBackward0"; }
void release_variables() override {
std::lock_guard<std::mutex> lock(mutex_);
self_.reset_data();
other_.reset_data();
}
SavedVariable self_;
at::ScalarType other_scalar_type;
at::ScalarType self_scalar_type;
SavedVariable other_;
};
variable_list MulBackward0::apply(variable_list&& grads) {
std::lock_guard<std::mutex> lock(mutex_);
IndexRangeGenerator gen;
auto self_ix = gen.range(1);
auto other_ix = gen.range(1);
variable_list grad_inputs(gen.size());
auto& grad = grads[0];
auto self = self_.unpack();
auto other = other_.unpack();
bool any_grad_defined = any_variable_defined(grads);
if (should_compute_output({ other_ix })) {
auto grad_result = any_grad_defined ? (mul_tensor_backward(grad, self, other_scalar_type)) : Tensor();
copy_range(grad_inputs, other_ix, grad_result);
}
if (should_compute_output({ self_ix })) {
auto grad_result = any_grad_defined ? (mul_tensor_backward(grad, other, self_scalar_type)) : Tensor();
copy_range(grad_inputs, self_ix, grad_result);
}
return grad_inputs;
}
我們去 tools/autograd/derivatives.yaml 中看看，看到就是使用了 mul_tensor_backward。
- name: mul.Tensor(Tensor self, Tensor other) -> Tensor
self: mul_tensor_backward(grad, other, self.scalar_type())
other: mul_tensor_backward(grad, self, other.scalar_type())
result: other_t * self_p + self_t * other_p
其最后也使用了 handle_r_to_c。
Tensor mul_tensor_backward(Tensor grad, Tensor other, ScalarType self_st) {
auto out = grad * other.conj();
return handle_r_to_c(self_st, out);
}
3.6 PermuteBackward
PermuteBackward 雖然沒有在上圖中體現，但是實際上存在，就是賦值操作。PermuteBackward 定義如下：
struct TORCH_API PermuteBackward : public Node {
using Node::Node;
variable_list apply(variable_list&& grads) override;
std::string name() const override { return "PermuteBackward"; }
void release_variables() override {
}
std::vector<int64_t> dims;
};
variable_list PermuteBackward::apply(variable_list&& grads) {
IndexRangeGenerator gen;
auto self_ix = gen.range(1);
variable_list grad_inputs(gen.size());
auto& grad = grads[0];
bool any_grad_defined = any_variable_defined(grads);
if (should_compute_output({ self_ix })) {
auto grad_result = any_grad_defined ? (permute_backwards(grad, dims)) : Tensor();
copy_range(grad_inputs, self_ix, grad_result);
}
return grad_inputs;
}
我們去 tools/autograd/derivatives.yaml 中看看，看到就是使用了 permute_backwards。
- name: permute(Tensor(a) self, int[] dims) -> Tensor(a)
self: permute_backwards(grad, dims)
result: auto_linear
permute_backwards 定義在 torch/csrc/autograd/FunctionsManual.cpp。
Tensor permute_backwards(const Tensor & grad, IntArrayRef fwd_dims) {
// invert the permutation
auto ndims = fwd_dims.size();
std::vector<int64_t> dims(ndims);
for(const auto i : c10::irange(ndims)) {
dims[at::maybe_wrap_dim(fwd_dims[i], ndims)] = i;
}
return grad.permute(dims);
}
我們接下來具體分析前向計算，看看其如何搭建依賴關系。
0x04 前向計算
因為篇幅所限，我們直接跳到 C++世界的核心之處。
4.1 減法實現
經過層層分發，減法最終調用到 torch/csrc/autograd/generated/VariableTypeEverything.cpp，PyTorch將會在這個函數中構建autograd的信息，其總體邏輯是：
- 1）減法操作會派發到某一個device之上，其中會進行前向計算結果Variable的構建。
- 2）派發到VariableType上時，會進行autograd信息的構建；
- 構建一個減法的反向計算函數 SubBackward0 實例，實例名字為 grad_fn。
- 設置反向計算時候使用的函數。
- 初始化SubBackward0實例的
next_edges_和其它相關成員，next_edges__成員的值來自前向傳播的輸入參數。- 如果輸入
Variable是leaf節點，則next_edges__ 來自輸入Variable的grad_accumulator_ - 如果 Varaible是非leaf節點，則
next_edges_來自Variable的grad_fn_。
- 如果輸入
- 使用步驟3中的Variable實例來初始化 SubBackward0 實例的
input_metadata_，
- 3）前向運算后得到新的Variable result，使用Variable::Impl進行構建。
- 4）設置計算歷史，使用步驟2) 中的 SubBackward0 實例 grad_fn 初始化該Variable實例的
autograd_meta_->grad_fn_成員。 - 5）返回 result。這里的 result 就是前向計算的結果，也就是我們示例之中的 Q。
具體代碼如下：
m.impl("sub.Tensor",
TORCH_FN(VariableType::sub_Tensor)
);
at::Tensor sub_Tensor(c10::DispatchKeySet ks, const at::Tensor & self, const at::Tensor & other, const at::Scalar & alpha) {
auto& self_ = unpack(self, "self", 0);
auto& other_ = unpack(other, "other", 1);
auto _any_requires_grad = compute_requires_grad( self, other );
(void)_any_requires_grad;
auto _any_has_forward_grad_result = isFwGradDefined(self) || isFwGradDefined(other);
(void)_any_has_forward_grad_result;
std::shared_ptr<SubBackward0> grad_fn; // 構建SubBackward0
if (_any_requires_grad) {
// 設置反向計算時候使用的函數
grad_fn = std::shared_ptr<SubBackward0>(new SubBackward0(), deleteNode);
// 設置下一條邊的所有輸入變量
grad_fn->set_next_edges(collect_next_edges( self, other ));
// 設置下一條邊的類型
grad_fn->other_scalar_type = other.scalar_type();
grad_fn->alpha = alpha;
grad_fn->self_scalar_type = self.scalar_type();
}
#ifndef NDEBUG
c10::optional<Storage> self__storage_saved =
self_.has_storage() ? c10::optional<Storage>(self_.storage()) : c10::nullopt;
c10::intrusive_ptr<TensorImpl> self__impl_saved;
if (self_.defined()) self__impl_saved = self_.getIntrusivePtr();
c10::optional<Storage> other__storage_saved =
other_.has_storage() ? c10::optional<Storage>(other_.storage()) : c10::nullopt;
c10::intrusive_ptr<TensorImpl> other__impl_saved;
if (other_.defined()) other__impl_saved = other_.getIntrusivePtr();
#endif
auto _tmp = ([&]() {
at::AutoDispatchBelowADInplaceOrView guard;
// 前向計算
return at::redispatch::sub(ks & c10::after_autograd_keyset, self_, other_, alpha);
})();
// 得到前向計算的輸出
auto result = std::move(_tmp);
if (grad_fn) {
// 將輸出variable與grad_fn綁定，grad_fn之中包含了計算梯度的function
// 設置計算歷史
set_history(flatten_tensor_args( result ), grad_fn);
}
if (_any_has_forward_grad_result) {
auto self_t_raw = toNonOptFwGrad(self);
auto self_t = self_t_raw.defined() ? self_t_raw : at::zeros_like(toNonOptTensor(self));
auto other_t_raw = toNonOptFwGrad(other);
auto other_t = other_t_raw.defined() ? other_t_raw : at::zeros_like(toNonOptTensor(other));
auto result_new_fw_grad = self_t - maybe_multiply(other_t, alpha);
if (result_new_fw_grad.defined()) {
// The hardcoded 0 here will need to be updated once we support multiple levels.
result._set_fw_grad(result_new_fw_grad, /* level */ 0, /* is_inplace_op */ false);
}
}
return result;
}
我們接下來逐一分析。首先要分析基礎函數，然后再回來分析 sub_Tensor。
4.3 邊基礎函數
我們首先介紹兩個構建邊相關的函數。
4.3.1 create_gradient_edge
create_gradient_edge代碼位於 torch/csrc/autograd/function.h。其作用是：
- 在給定的"變量"和"函數"之間創建一個"邊"，該函數是該變量的梯度函數（即，在后向傳播過程中計算該變量梯度的函數）。
- 此函數將設置"variable"的"grad_fn"屬性。
create_gradient_edge 方法假定'Variable'是梯度函數的新輸入，因此其'input_nr'等於function->num_inputs()。此外，它還將"節點"的輸入數增加一。
如果不希望增加"節點"的"num_inputs"，請直接使用"set_gradient_edge"。從功能上來說，create_gradient_edge 大約相當於 variable.set_gradient_edge(function, function->add_input_metadata(variable.dispatch_type(), variable.sizes()))。
/// Create an `Edge` between the given `variable` and the `function`, which is
/// assumed to be the gradient function of this variable (i.e. the function
/// through which this variable is backpropagated during the backward pass).
/// This sets the `grad_fn` property of the `variable`. This function assumes
/// that the `Variable` is a new input to the gradient function and its
/// `input_nr` thus equal to `function->num_inputs()`. Additionally, it
/// increments the `Node`'s number of inputs by one. Approximately
/// equivalent to `variable.set_gradient_edge(function,
/// function->add_input_metadata(variable.dispatch_type(), variable.sizes()))`.
/// If you don't want the `Node`'s `num_inputs` to be incremented, use
/// `set_gradient_edge` directly.
inline void create_gradient_edge(
Variable& variable,
std::shared_ptr<Node> function) {
// Copy before move.
const auto input_nr = function->add_input_metadata(variable);
impl::set_gradient_edge(variable, {std::move(function), input_nr});
}
4.3.2 set_gradient_edge
set_gradient_edge 代碼位於 torch/csrc/autograd/variable.cpp。
配置歷史的操作會最終調用到這里，這是使用 edge 來真正配置了本張量如何計算梯度，而且是配置到了 Variable 類之上的 autograd_meta_。即獲取 Tensor 的 autograd_meta_，配置其 grad_fn_ 和 output_nr_。
void set_gradient_edge(const Variable& self, Edge edge) {
auto* meta = materialize_autograd_meta(self);
meta->grad_fn_ = std::move(edge.function); // 配置梯度函數
meta->output_nr_ = edge.input_nr; // 配置梯度函數的第幾個輸出
// For views, make sure this new grad_fn_ is not overwritten unless it is necessary
// in the VariableHooks::grad_fn below.
// This logic is only relevant for custom autograd Functions for which multiple
// operations can happen on a given Tensor before its gradient edge is set when
// exiting the custom Function.
auto diff_view_meta = get_view_autograd_meta(self);
if (diff_view_meta && diff_view_meta->has_bw_view()) {
diff_view_meta->set_attr_version(self._version());
}
}
其中，materialize_autograd_meta 代碼如下，其作用就是從 Tensor 之中獲取 autograd_meta_。
AutogradMeta* materialize_autograd_meta(const Variable& self) {
TORCH_CHECK(self.defined(), "cannot call materialize_autograd_meta() on undefined tensor");
auto p = self.unsafeGetTensorImpl();
if (!p->autograd_meta()) {
p->set_autograd_meta(std::make_unique<AutogradMeta>());
}
return get_autograd_meta(self);
}
get_view_autograd_meta 代碼如下，返回了 DifferentiableViewMeta。
DifferentiableViewMeta* get_view_autograd_meta(const Variable& self) {
// NB: return nullptr if self is not a view
AutogradMeta* meta = get_autograd_meta(self);
if (meta && meta->is_view_) {
return static_cast<DifferentiableViewMeta*>(meta);
} else {
return nullptr;
}
}
4.4 構建網絡
我們已經分析了 SubBackward0 和 基礎函數，接下返回來分析 sub_Tensor 的實現。首先是構建后向傳播網絡。
- 首先，構建一個 SubBackward0 grad_fn。
- 其次，對 grad_fn 進行設置，主要是 使用collect_next_edges()搜集 sub 操作兩個變量的，然后進行set_next_edges。
- 然后，進行前向計算，得到前向計算的輸出。
- 最后，將輸出variable加入到history之中，將輸出variable與grad_fn綁定。
下面代碼只是保留 sub_Tensor 關鍵部分。
std::shared_ptr<SubBackward0> grad_fn;
if (_any_requires_grad) {
// 反向計算時候使用的函數
grad_fn = std::shared_ptr<SubBackward0>(new SubBackward0(), deleteNode);
// 設置下一條邊的所有輸入變量
grad_fn->set_next_edges(collect_next_edges( self, other ));
grad_fn->other_scalar_type = other.scalar_type();
grad_fn->alpha = alpha;
grad_fn->self_scalar_type = self.scalar_type();
}
auto _tmp = ([&]() {
at::AutoDispatchBelowADInplaceOrView guard;
// 前向計算
return at::redispatch::sub(ks & c10::after_autograd_keyset, self_, other_, alpha);
})();
// 得到前向計算的輸出
auto result = std::move(_tmp);
if (grad_fn) {
// 將輸出variable與grad_fn綁定，grad_fn之中包含了計算梯度的function
// 將本身計算加入到計算歷史之中
set_history(flatten_tensor_args( result ), grad_fn);
}
4.5 構建邊
構建網絡的關鍵部分就是構建邊，這里是配置反向傳播的輸出邊（輸出邊對應了SubBackward0的兩個輸入），其中有兩步驟:
- 使用 collect_next_edges 來收集輸入參數（張量）的邊，得到了后續邊，后續邊就是兩個輸入參數 self和other的gradient_edge()。
- 使用 set_next_edges 把邊配置到張量上。當set_next_edges調用完成后，一個 Node 的 next_edges_成員（類型為std::vector
）就會初始化完成。
4.5.1 獲取邊
collect_next_edges 函數就是用來根據輸入變量來獲取邊。其實，collect_next_edges 就是得到 self 和 other 的gradient_edge。
4.5.1.1 gradient_edge
gradient_edge方法作用是返回通過Variable的 grad_fn_構建的Edge實例，邏輯如下：
- 就是如果一個節點有 grad_fn：
- 說明節點是內部節點（通過運算內部創建的）。
- grad_fn_就是這個Variable的gradient function，
- 那么就使用 grad_fn來構建一個 Edge返回。
- 如果一個節點沒有 grad_fn：
- 說明是葉子節點（用戶創建的）。
- grad_fn_ 是這個Variable的gradient accumulator，也就是一個AccumulateGrad類（Function子類）的實例。PyTorch 使用grad_accumulator來累加輸出給這個Variable的梯度。
- 使用grad_accumulator來構建一個 Edge返回。
代碼如下，需要注意的是，output_nr是當前variable在前向計算時是第幾個輸出，對於單輸出的算子比如add或者mul來說，output_nr一般都是0，但對於多輸出的算子比如split，則output_nr可能是0,1,2...。
Edge gradient_edge(const Variable& self) {
// If grad_fn is null (as is the case for a leaf node), we instead
// interpret the gradient function to be a gradient accumulator, which will
// accumulate its inputs into the grad property of the variable. These
// nodes get suppressed in some situations, see "suppress gradient
// accumulation" below. Note that only variables which have `requires_grad =
// True` can have gradient accumulators.
// self.grad_fn() 這里觸發了一個調用，得到了一個SubBackward0實例
if (const auto& gradient = self.grad_fn()) { // 這是一個中間節點，gradient 是一個Function
return Edge(gradient, self.output_nr()); // self.output_nr() 表示本Edge是function的第n個輸入。前向傳播時候的第 n 個輸出在反向傳播時候就是第 n 個輸入。
} else {
return Edge(grad_accumulator(self), 0); // 這是一個葉子節點，所以生成一個AccumulateGrad，0表示本Edge是function的第一個輸入
}
}
4.5.1.2 gradient accumulator
這里有一步需要注意，就是 gradient_edge 方法中，有這樣一個語句 return Edge(grad_accumulator(self), 0)，這個代碼實際是觸發Variable::grad_accumulator()調用。
在一個Variable第一次調用這個API的時候，會生成一個AccumulateGrad 來初始化它的 grad_accumulator_成員，代碼如下：
std::shared_ptr<Node> grad_accumulator(const Variable& self) {
auto autograd_meta = get_autograd_meta(self);
if (!autograd_meta) {
return nullptr;
}
if (autograd_meta->grad_fn_) {
throw std::logic_error(
"grad_accumulator() should be only called on leaf Variables");
}
if (!autograd_meta->requires_grad_) {
return nullptr;
}
std::lock_guard<std::mutex> lock(autograd_meta->mutex_);
auto result = autograd_meta->grad_accumulator_.lock();
if (result)
return result;
c10::raw::intrusive_ptr::incref(self.unsafeGetTensorImpl());
auto intrusive_from_this = c10::intrusive_ptr<at::TensorImpl>::reclaim(self.unsafeGetTensorImpl());
// 這里會初始化一個AccumulateGrad，配置給grad_accumulator_
result = std::make_shared<AccumulateGrad>(Variable(std::move(intrusive_from_this)));
autograd_meta->grad_accumulator_ = result;
return result;
}
4.5.1.3 AccumulateGrad
AccumulateGrad 定義位於 torch/csrc/autograd/functions/accumulate_grad.h
struct TORCH_API AccumulateGrad : public Node {
explicit AccumulateGrad(Variable variable_); // 必須用一個Variable構建
variable_list apply(variable_list&& grads) override; // 接收一個list的Variable的實例
Variable variable;
};
其構造函數在 torch/csrc/autograd/functions/accumulate_grad.cpp。
這會new一個AccumulateGrad對象，使用UINT64_MAX 來初始化Function的sequence_nr_成員。
AccumulateGrad::AccumulateGrad(Variable variable_)
: Node(/*sequence_nr=*/UINT64_MAX),
variable(std::move(variable_)) {
add_input_metadata(variable);
}
4.5.1.4 收集邊
collect_next_edges 這里建立了邊。收集了所有輸入的邊。
/// Return the next edges of all the given variables, or tuples of variables.
template <typename... Variables>
edge_list collect_next_edges(Variables&&... variables) {
detail::MakeNextFunctionList make; // 這里將調用gradient_edge
// next_edges_成員的值來自前向時候的輸入參數
make.apply(std::forward<Variables>(variables)...);
return std::move(make.next_edges);
}
MakeNextFunctionList 的定義如下，apply 時候會構建 gradient_edge，這就對應了前面所說的 gradient_edge 等小節。
struct MakeNextFunctionList : IterArgs<MakeNextFunctionList> {
edge_list next_edges;
using IterArgs<MakeNextFunctionList>::operator();
void operator()(const Variable& variable) {
if (variable.defined()) {
next_edges.push_back(impl::gradient_edge(variable)); // 調用gradient_edge
} else {
next_edges.emplace_back();
}
}
void operator()(const c10::optional<Variable>& variable) {
if (variable.has_value() && variable->defined()) {
next_edges.push_back(impl::gradient_edge(*variable)); // 調用gradient_edge
} else {
next_edges.emplace_back();
}
}
};
此時得到了 edge_list，但是沒有和 SubBackward0 建立聯系。
+------------------------+ +----------------------+
| SubBackward0 | | |
| | | Compute the gradient |
| apply +-----------------> | |
| | +----------------------+
| |
| |
| next_edges_ |
| |
| other_scalar_type |
| |
| alpha |
| |
| self_scalar_type |
| |
| input_metadata_ |
| |
+------------------------+
+-----------------------------------------------------+
| edge_list |
| |
| [(MulBackward0(self), 0), (PowBackward0(other), 0)] |
| |
+-----------------------------------------------------+
4.5.2 配置邊
獲取到了所有輸出邊之后，接下來就要設置到 SubBackward0 的 next_edges_ 之上，一定要注意，next_edges_成員的值來自前向傳播時候的輸入參數。
void set_next_edges(edge_list&& next_edges) {
next_edges_ = std::move(next_edges); // 這里設置了邊
for(const auto& next_edge : next_edges_) {
update_topological_nr(next_edge);
}
}
update_topological_nr 會依據輸出邊來設置 topological_nr
void update_topological_nr(const Edge& edge) {
Node* node = edge.function.get();
if (node) {
auto topo_nr = node->topological_nr();
if (topological_nr_ <= topo_nr) {
topological_nr_ = topo_nr + 1;
}
}
}
結合我們的例子，此時應該如下圖，下圖中 0 的意義舉例如下：(PowBackward0(other), 0) 中的 0 表示SubBackward0 的計算輸出是 PowBackward0 的第一個輸入（原始冪運算只有一個輸出）。
+------------------------+ +----------------------+
| SubBackward0 | | |
| | | Compute the gradient |
| apply +-----------------> | |
| | +----------------------+
| |
| | +-----------------------------------------------------+
| next_edges_ +-----------> | edge_list |
| | | |
| other_scalar_type | | [(MulBackward0(self), 0), (PowBackward0(other), 0)] |
| | | |
| alpha | +-----------------------------------------------------+
| |
| self_scalar_type |
| |
| input_metadata_ |
| |
+------------------------+
4.6 配置歷史
接下來是配置歷史，result 是之前代碼計算出來的前向傳播輸出，這里其實是配置反向傳播的輸入參數 和 輸入如何計算。
if (grad_fn) { // grad_fn 就是 std::shared_ptr<SubBackward0>
// 將輸出variable與grad_fn綁定，grad_fn之中包含了計算梯度的function
set_history(flatten_tensor_args( result ), grad_fn);
}
set_history 會把前向傳播結果加入到history之中，具體就是遍歷結果中的張量，然后把每一個張量加入到history。其中關鍵一點是調用了前面提到的 set_gradient_edge，把 grad_fn（就是 SubBackward0）配置給了result.autograd_meta_ 的 grad_fn_。
回憶一下 Tensor 的成員變量 grad_fn 定義。
grad_fn：指向一個Function對象。
- 這個Function對象用來在反向傳播時候計算輸入的梯度。
- 若本張量是非葉節點，則 Function 是向葉節點方向操作的反向傳播函數，比如例子里 O 節點對應的函數就是MulBackward，即乘法操作的反向函數；
經過對比，就可以知道，前向操作的輸入 result 在反向傳播計算梯度時候，就會使用 grad_fn_ 來計算梯度，就是我們這里的 SubBackward0。這樣就設置了反向傳播如何針對輸入來計算梯度。
具體 set_history 代碼如下：
inline void set_history(
at::Tensor& variable,
const std::shared_ptr<Node>& grad_fn) {
if (variable.defined()) {
// grad_fn 的 input_metadata 之中添加了輸出實例，輸出實例在反向傳播時候就是輸入
auto output_nr = grad_fn->add_input_metadata(variable);
// 輸出實例 result 中設置上了grad_fn，這里配置了邊，邊就是 {grad_fn, output_nr}。
// output_nr_被賦值成了"當前Variable信息在input_metadata_中的index"。
impl::set_gradient_edge(variable, {grad_fn, output_nr});
} else {
// 設置成未定義
grad_fn->add_input_metadata(Node::undefined_input());
}
}
inline void set_history(
std::vector<Variable>&& variables,
const std::shared_ptr<Node>& grad_fn) {
for (auto& variable : variables) {
set_history(variable, grad_fn); // 調用到上面的函數
}
}
4.6.1 配置meta
配置歷史中，首先是配置input_metadata。將 input_metadata 之中添加了輸出實例 result，輸出實例 result 在反向傳播時候就是輸入。
4.6.1.1 input_metadata_
Node 類之中，input_metadata_ 的類型如下：
at::SmallVector<InputMetadata, 2> input_metadata_;
具體 InputMetadata 定義如下：
struct InputMetadata {
InputMetadata(const at::TensorOptions options, at::IntArrayRef shape, at::Device device)
: options_{options}, shape_{shape}, device_{device} {
stream_ = c10::impl::getDeviceGuardImpl(device_.type())->getStream(device_);
}
InputMetadata(const at::Tensor& t)
: InputMetadata(t.options(), t.sizes(), t.device()) { }
private:
const at::TensorOptions options_;
at::DimVector shape_;
at::Device device_ = at::kCPU;
c10::Stream stream_ = c10::Stream(c10::Stream::Default::DEFAULT, device_);
};
4.6.1.2 配置meta
add_input_metadata 方法之中 將meta 信息配置如下：
/// Adds the type and shape metadata for a new input. Returns the index of
/// of the new input.
uint32_t add_input_metadata (
const at::TensorOptions& options
, at::IntArrayRef shape
, at::Device device) noexcept {
uint32_t input_nr = input_metadata_.size();
input_metadata_.emplace_back(options, shape, device);
return input_nr;
}
配置之后，input_metadata_ 里面就增加了一個新 InputMetadata，InputMetadata 內容就是 輸出變量 result 的部分信息 (type, shape, device)，input_metadata_ 中的 index 就是 AutogradMeta 之中的 output_nr_。
所以，此時內存大致如下：
+-------------------------------------------------------------------------------------------------------------+
self +--+ | sub_Tensor |
| | +--------------------------+ +----------------------+ |
+---->+ |SubBackward0 | | | |
| | | | | Compute the gradient | |
other +--+ | +--> grad_fn---> | apply +-----------------> | | |
| | | | +----------------------+ |
| | | | |
| | | | +-----------------------------------------------------+ |
| | | next_edges_ +-----------> | edge_list | |
| | | | | | |
| | | other_scalar_type | | [(PowBackward0(self), 0), (PowBackward0(other), 0)] | |
| | | | | | |
| | | alpha | +-----------------------------------------------------+ |
| | | | |
| | | self_scalar_type | +------------------------------------------------------+ |
| | | | | | |
| | | input_metadata_ +-------> | [(type of result, shape of result, device of result)]| |
| | | | | | |
| | +--------------------------+ +------------------------------------------------------+ |
| | |
| | |
| | +-----------------------+ +---------------------------------------+ |
| | |result | | DifferentiableViewMeta | |
| | | | | | |
| | | autograd_meta_ +-----------> | grad_ grad_accumulator_ | |
| | | | | | |
| | +-----------------------+ | | |
| +--------------------------------------------------------- grad_fn_ output_nr_ | |
| | | |
| +---------------------------------------+ |
+-------------------------------------------------------------------------------------------------------------+
手機如下：
4.7 印證
我們和之前示例對照印證，把示例代碼繼續細化，得到：
a = torch.tensor(2., requires_grad=True)
b = torch.tensor(6., requires_grad=True)
X = a ** 3
Y = 3 * X
Z = b ** 2
Q = X - Z
external_grad = torch.tensor(1.)
Q.backward(gradient=external_grad)
print(a.grad)
print(b.grad)
看看運行時變量如下，因為 Q = X - Z 是減法，所以對應的反向操作就是 SubBackward0：
Q = {Tensor} tensor(-28., grad_fn=<SubBackward0>)
X = {Tensor} tensor(8., grad_fn=<PowBackward0>)
Y = {Tensor} tensor(24., grad_fn=<MulBackward0>)
Z = {Tensor} tensor(36., grad_fn=<PowBackward0>)
a = {Tensor} tensor(2., requires_grad=True)
b = {Tensor} tensor(6., requires_grad=True)
我們再具體看看，注意，(<PowBackward0 object at 0x00000177300F4688>, 0) 這里的 0 表示本Node是PowBackward0的第0的輸出，也就是唯一輸出。
Q = {Tensor}
grad_fn = {SubBackward0}
next_functions = {tuple: 2}
0 = {tuple: 2} (<PowBackward0 object at 0x00000177300F4688>, 0)
1 = {tuple: 2} (<PowBackward0 object at 0x00000177300F46C8>, 0)
X = {Tensor}
grad_fn = {PowBackward0}
next_functions = {tuple: 1}
0 = {tuple: 2} (<AccumulateGrad object at 0x00000177300F49C8>, 0)
Z = {Tensor}
grad_fn = {PowBackward0}
next_functions = {tuple: 1}
0 = {tuple: 2} (<AccumulateGrad object at 0x00000177301003C8>, 0)
Y = {Tensor}
grad_fn = {MulBackward0}
next_functions = {tuple: 2}
0 = {tuple: 2} (<PowBackward0 object at 0x0000017730100CC8>, 0)
1 = {tuple: 2} (None, 0)
對應簡要圖是：
對應邏輯：
-
- 以 self 和 other 兩個張量為參數，調用 sub_Tensor
-
- 使用 grad_fn = std::shared_ptr
(new SubBackward0(), deleteNode); 構建一個SubBackward0。其中，grad_fn 的 next_edges_成員的值來自前向傳播的輸入參數，就是self 和 other。
- 使用 grad_fn = std::shared_ptr
-
- 使用 at::redispatch::sub 進行前向計算，得到 result。
-
-
使用 set_history 設置計算歷史。set_history 這里包含兩個部分
-
使用 output_nr = grad_fn->add_input_metadata(variable) 為 grad_fn 的 input_metadata 之中添加了輸出實例。
-
使用 impl::set_gradient_edge(variable, {grad_fn, output_nr}) 給 輸出實例 result 的屬性
autograd_meta_->grad_fn_中設置上了grad_fn。
-
-
- 最后返回了 result。
可以看到，sub_Tensor 針對 result 做了如下配置：
- 如何知道調用反向計算 ：result 就是前向計算的結果，result 之中有
autograd_meta_，其是一個 DifferentiableViewMeta 類型，DifferentiableViewMeta 的 grad_ 和 grad_fn_ 就是反向計算的梯度函數。grad_fn_ 指向了 SubBackward0。 - 反向傳播如何計算 ：調用 SubBackward0 計算。
- SubBackward0 的輸入 ：得到了前向計算的輸出 result（其會在反向傳播時候作為輸入變量，就是設定到了 SubBackward0.input_metadata_ 之上）。
- SubBackward0 的輸出 ：構建了
next_edges_作為其反向傳播時候的輸出邊。根據next_edges_就能得到反向傳導圖了。
其邏輯圖如下：
+---------------------------------------------------------------------------------------------------------------+
self +--+ | sub_Tensor +--------------------------+ +----------------------+ |
| | |SubBackward0 | | | |
+---->+ 2 | | | Compute the gradient | |
| 1 | +-----> grad_fn +-----> | apply +-----------------> | | |
other +--+ | | | | +----------------------+ |
| | | | |
| | | | +----------------------+ |
| | | next_edges_ +-----------> | edge_list | |
| | | | | | |
| | | other_scalar_type | | self, other | |
| | | | | | |
| | | alpha | +----------------------+ |
| | | | |
| | | self_scalar_type | |
| | | | |
| | | input_metadata_ +------> [result] |
| | | | ^ |
| | +--------------------------+ | |
| | | 5 |
| | | |
| | 3 result = at::redispatch::sub +--------------------------------------------------------+ |
| | | | | |
| | | + | |
| | | output_nr = grad_fn+>add_input_metadata(variable) | |
| | 4 set_history(result, grad_fn) +-------> | | |
| | | impl::set_gradient_edge(variable,a{grad_fn, output_nr})| |
| | | + | |
| +----------------------------+ | | | |
| 6 | +--------------------------------------------------------+ |
| | | |
| +-----------------------+ | +-----------------------------------+ | 7 |
| |result | | | DifferentiableViewMeta | | |
| | | | | | <---+ |
| | autograd_meta_ +---------------->+ | |
| | | | | grad_ grad_accumulator_ | |
| | | | | | |
| | | +--------+grad_fn_ output_nr_ | |
| | | | | |
| +------------+----------+ +-----------------------------------+ |
| | |
+---------------------------------------------------------------------------------------------------------------+
|
result | 7
v
手機如下：
至此，前向計算分析完成，我們下一篇開始介紹后向傳播。
0xFF 參考
https://github.com/KeithYin/read-pytorch-source-code/
pytorch學習筆記（十三）：backward過程的底層實現解析
How autograd encodes the history
https://pytorch.org/tutorials/beginner/blitz/autograd_tutorial.html