Pytorch1.4 ATEN/native 代碼 淺析

Note：近日閱讀Pytorch1.4代碼，發現與1.0有些差異，而未能找到關於native這些算子的代碼解析文章，於是便自己硬看。由於是初次接觸，若有錯誤，還請指出。

ATEN/native 簡介

代碼路徑：
pytorch/aten/src/ATEN/native

區別於 aten/THNN （即TH系列）的庫，引用官方文檔的話，native是：

ATen "native" functions are the modern mechanism for adding operators and functions to ATen (they are "native" in contrast to legacy functions, which are bound via TH/THC cwrap metadata). Native functions are declared in native_functions.yaml and have implementations defined in one of the cpp files in this directory.

即： ATEN/native 下的文件是使用C++實現的operaters(或者說"層")

特點

• 有C++和python API
  • 在C++中，namespace是 at::

關於自定義operator：

添加operator：

查閱 native 文件夾下的 Readme

自定義operator的自動微分

• 如果是使用其它（支持自動微分的）op拼起來的算子，則不需要實現backward函數
• 如果是自己實現的代碼，需要實現 foo_backward 並添加至tools/autograd/derivatives.yaml

以 activation 為例：

native/Activation.cpp

我理解這是一個‘入口’，將activation的各個實現包裝起來

• hardtanh
• elu & celu & selu & rrelu & prelu
• softplus
• threshold
• hardshrink & softshrink
• 等等

以及上面算子的不同操作形式變體（inplace與否）

以 elu 為例

native/Activation.cpp定義了如下接口：

• elu : 前向
  • elu_ & elu_out ：inplace elu
• elu_backward : 反向
  • elu_backward_out : inplace

Tensor elu(
    const Tensor& self,
    Scalar alpha,
    Scalar scale,
    Scalar input_scale) {
  Tensor result;
  auto iter = TensorIterator::unary_op(result, self);
  elu_stub(iter.device_type(), iter, alpha, scale, input_scale);
  return iter.output();
}

前向代碼調用的是 elu_stub

定義：

• native/cpu/Activation.cpp: REGISTER_DISPATCH(elu_stub, &elu_kernel);
• native/cuda/Activation.cu : REGISTER_DISPATCH(elu_stub, &elu_kernel);

反向代碼也類似，調用注冊的 _backward_stub

Tensor elu_backward(
    const Tensor& grad_output,
    Scalar alpha,
    Scalar scale,
    Scalar input_scale,
    const Tensor& output) {
  Tensor result;
  auto iter = TensorIterator::binary_op(result, grad_output, output);
  elu_backward_stub(iter.device_type(), iter, alpha, scale, input_scale);
  return iter.output();
}

native/cpu/Activation.cpp

該文件是激活函數的cpu代碼實現：

void elu_kernel(TensorIterator& it, Scalar alpha, Scalar scale, Scalar input_scale) {
  AT_DISPATCH_FLOATING_TYPES(it.dtype(), "elu_cpu", [&]() {
    auto negcoef = alpha.to<scalar_t>() * scale.to<scalar_t>();
    auto poscoef = scale.to<scalar_t>();
    auto negiptcoef = input_scale.to<scalar_t>();
    cpu_kernel(it, [=](scalar_t a) -> scalar_t {
      return a <= scalar_t(0) ? (std::exp(a * negiptcoef) - scalar_t(1)) * negcoef : a * poscoef;
    });
  });
}

前向：

• x > 0 : x*scale
• x <= 0 : negcoef*(e^(x*neginputcoef) - 1)

void elu_backward_kernel(TensorIterator& it, Scalar alpha, Scalar scale, Scalar input_scale) {
  AT_DISPATCH_FLOATING_TYPES(it.dtype(), "elu_backward_cpu", [&]() {
    auto negcoef = alpha.to<scalar_t>() * scale.to<scalar_t>();
    auto poscoef = scale.to<scalar_t>();
    auto negiptcoef = input_scale.to<scalar_t>();
    cpu_kernel(it, [=](scalar_t a, scalar_t b) -> scalar_t {
      return b <= scalar_t(0) ? a * negiptcoef * (b + negcoef) : a * poscoef;
    });
  });
}

這里有a,b兩個參數，分別應該是 grad_output, output

elu的梯度：

• output <= 0 : grad_output*neginputcoef*output
• output > 0 : output*scale