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"]

  DISTRIBUTED RPC FRAMEWORK APIS (NOW STABLE)

  分布式 RPC 框架在 1.4 版本中以实验性形式出现,现在它已经处于稳定状态。此过程涉及许多增强功能和 bug 修复,以使分布式 RPC 框架总体上更可靠和更健壮。同时还添加了两个新特性,包括概要分析支持,在 RPC 中使用 TorchScript 函数以及一些易于使用的增强功能。

  此外,从 1.5 开始,PyTorch 不再支持 Python 2,今后,对 Python 的支持将仅限于 Python 3,特别是 Python 3.5、3.6、3.7 和 3.8。

  更多具体的细节见发布公告:

  https://pytorch.org/blog/pytorch-1-dot-5-released-with-new-and-updated-apis