PyTorch 1.5 发布，C++ 前端重大更新、引入新存储格式提高性能

PyTorch 1.5 发布，C++ 前端重大更新、引入新存储格式提高性能-风君雪科技博客

　　PyTorch 1.5 发布了，此版本包括新增加了几个主要的 API 和一些改进、对 C++ 前端有重大更新、用于计算机视觉模型的“通道维持（channels last）”存储格式，以及用于模型并行训练的分布式 RPC 框架的稳定版本。

　　同时，此版本还提供了针对 hessian 和 jacobian 的 autograd 的新 API，以及一个受 pybind 启发的允许创建 Custom C++ Class（定制类）的 API。

　　C++ FRONTEND API (STABLE)

　　之前被标记为实验性质的 C++ 前端 API 现在已经与 Python 地位相当，并且总体特性已移至“稳定”状态。一些主要亮点包括：

现在覆盖率达到约 100％，并提供有关 C++ torch::nn 模块/函数的文档，开发者可以轻松地将其模型从 Python API 转换为 C++ API。
C++ 中的优化器与 Python 中的有所不同：C++ 优化器不能将参数组作为输入，而 Python 可以。此外，步骤函数的实现也不完全相同。在 1.5 版本中，C++ 优化器的行为将始终与 Python 等效器相同。
C++ 中缺少张量多维度索引 API 是一个众所周知的问题，并在 PyTorch GitHub issue 跟踪器和论坛中引起了很多讨论。以前的解决方法是结合使用 arrow/select/index_select/masked_select，与 Python API 的优雅张量 [:, 0, …， mask] 语法相比，它比较笨拙且容易出错。在 1.5 版本中，开发者可以使用 tensor.index ({Slice ()， 0, “…”， mask}) 达到相同的目的。

　　‘CHANNELS LAST’ MEMORY FORMAT FOR COMPUTER VISION MODELS (EXPERIMENTAL)

　　用于计算机视觉模型的“通道维持（channels last）”存储格式，目前处于实验阶段。该格式的内存布局可以充分发挥性能高效的卷积算法和硬件（NVIDIA 的 Tensor Core、FBGEMM、QNNPACK）的能力。此外，它被设计为通过运算符自动传播，从而允许在内存布局之间轻松切换。

　　CUSTOM C++ CLASSES (EXPERIMENTAL)

　　此版本添加了一个新的 API torch.CutomClassHolder，用于将自定义 C++ 类同时绑定到 TorchScript 和 Python，该 API 的语法与 pybind11 几乎相同。它允许开发者向 TorchScript 类型系统和运行时系统暴露其 C++ 类及其方法，以便他们可以实例化和操作 TorchScript/Python 中的任意 C++ 对象。C++ 绑定示例：

template <class T>
struct MyStackClass : torch::CustomClassHolder {
  std::vector<T> stack_;
  MyStackClass (std::vector<T> init) : stack_(std::move (init)) {}

  void push (T x) {
    stack_.push_back (x);
  }
  T pop () {
    auto val = stack_.back ();
    stack_.pop_back ();
    return val;
  }
};

static auto testStack =
  torch::class_<MyStackClass<std::string>>("myclasses", "MyStackClass")
      .def (torch::init<std::vector<std::string>>())
      .def ("push", &MyStackClass<std::string>::push)
      .def ("pop", &MyStackClass<std::string>::pop)
      .def ("size", [](const c10::intrusive_ptr<MyStackClass>& self) {
        return self->stack_.size ();
      });

　　它暴露了可以在 Python 和 TorchScript 中使用的类，如下所示：

@torch.jit.script
def do_stacks (s : torch.classes.myclasses.MyStackClass):
    s2 = torch.classes.myclasses.MyStackClass (["hi", "mom"])
    print (s2.pop ()) # "mom"
    s2.push ("foobar")
    return s2 # ["hi", "foobar"]