pytorch学习笔记（十三）：backward过程的底层实现解析

pytorch

python

C/C++

pytorch

pytorch

pytorch 是如何构建反向传导图的

pytorch

pytorch

pytorch

pytorch

// Add 函数 计算方法的实现， 构建反向传导图的关键在 wrap_outputs
auto Add::apply(const variable_list& inputs) -> variable_list {
check_input_variables("Add", inputs, 2);
auto& input1 = inputs[0]->data;
auto& input2 = inputs[1]->data;
AutoGPU guard(input1->getDevice());

bool first_sparse = input1->isSparse();
auto output = first_sparse ? input2->newTensor() : input1->newTensor();
if (first_sparse) {
output->cadd(*input2, *input1);
} else {
output->cadd(*input1, *input2);
}

return wrap_outputs(inputs, as_tensor_list(std::move(output)), [&](FunctionFlags f) {
return std::make_shared<AddBackward>(std::move(f));
});
};

wrap_outputs

variable_list wrap_outputs(const variable_list& inputs, tensor_list&& outputs,
function_constructor ctr) {

// 使用 inputs variables 来计算 反向传导的 Function 的 flag (f.is_executable, f.is_volatile)和 next_functions
// 这里需要搞清楚的一点是：inputs 是前向传导的 inputs，从它可获得的信息有：当前 函数的 反向传导函数 是否 可执行，

auto flags = Function::flags(inputs);
variable_list result;

// 开始创建 返回的 Variable 了。
result.reserve(outputs.size());

if (flags.is_volatile) {  // 如果 is_volatile=true, 那么输出的 Variable 的 is_volatile=true
for (auto& output : outputs) {
if (output.defined()) { // 因为 可能返回 None 嘛，所以这里 check 一下
result.emplace_back(make_variable(output, false, true)); // requires_grad=false, is_volatile=true
} else {
result.emplace_back();
}
}
} else {  // 如果 volatile=false， 难道也不管 is_executable 了吗？
// ctr 是一个 lambda 函数， 它返回一个 std::shared_ptr<GradFn>
// 梯度 使用 Function::flags 计算出来的 flags 其实是给 Backward 用的。
auto grad_fn = ctr(std::move(flags));  // 用 flags(is_executable, is_volatile) 创建出来一个 Function。
for (auto& output : outputs) {
if (output.defined()) {
result.emplace_back(make_variable(output, grad_fn));
} else {
// forward 的输出 变量个数，就是 backward 的输入变量个数。
++grad_fn->num_inputs;
result.emplace_back();
}
}
}
return result;
}

pytorch 的反向传导计算过程

a = o+e
c = a+b
d = a+e
res = c+d

res

o

a

o

a

d

e

backward

a

pytorch

function

输入

pytorch

FunctionTask

InputBuffer

grad_fn(a,o,e)

a

grad_fn(a,o,e)

FunctionTask

ready_queue

backward

struct FunctionTask {
// 每个 FunctionTask 中都维护着一个 base GraphTask
GraphTask* base;
std::shared_ptr<Function> fn; // 代表 grad_fn
InputBuffer inputs; // 累积 grad_fn 的输入

FunctionTask(GraphTask* base, std::shared_ptr<Function> fn, InputBuffer inputs)
: base(base)
, fn(fn)
, inputs(std::move(inputs)) {}
};

struct ReadyQueue {
// 用来 存放可被 执行的 FunctionTask
// queue 是 FunctionTask 的 一个 双端队列
std::deque<FunctionTask> queue;
// std::condition_variable 条件变量，同步的时候会用到
// 用 unique_lock (over mutex) 来进行操作
std::condition_variable not_empty;
std::mutex mutex;

void push_front(FunctionTask item);
FunctionTask pop_back();
};

struct GraphTask {
// 记录整个反向计算图的依赖情况 等等。
std::exception_ptr exception;
// Indicates if an error occurred while executing any task.  When this is
// true, it signals all threads to stop executing.
std::atomic_bool has_error;

// 剩余 tasks。 在 ReadyQueue 的 push方法 中加一， 在 evaluate_function 中减一操作
std::atomic<uint64_t> outstanding_tasks;
bool keep_graph;
bool has_any_work;
// 用来 给 notify_all 加锁的
std::mutex mutex;
// Notified when a task finishes executing.  Check outstanding_tasks to see
// if all tasks are done.
std::condition_variable not_done;
const Engine::pre_callback_map& pre_callbacks;
const Engine::post_callback_map& post_callbacks;

//用来存放 没有 准备好的 FunctionTask
std::unordered_map<Function*, InputBuffer> not_ready;
// 记录 所有 Function 节点的 依赖
std::unordered_map<Function*, int> dependencies;

// 这个来 表示 GraphTask 是在哪个 device 上创建的
int owner;

};

struct Engine{
// 反向传导计算引擎
}