Tensor Comprehensions將數(shù)學符號快速轉換為高性能機器學習代碼

Tensor Comprehensions 是 Facebook AI  研究院開源的  C++  庫及數(shù)學語言，功能齊全，能有效填補研究人員于數(shù)學運算領域的溝通鴻溝，并基于各種硬件后端上大規(guī)模運行工程模型。

Tensor Comprehensions 采用了 Just-In-Time 的編譯自動生成機器學習社區(qū)所需的高性能代碼，并被設計為高度可移植的。通過 Tensor Comprehensions，研究人員能夠以數(shù)學符號的方式進行編寫，系統(tǒng)能夠根據(jù)需求進行編譯調整，并輸出專業(yè)的代碼。

示例：

  #include <ATen/ATen.h>
  #include "tc/aten/aten_compiler.h"
  #include "tc/core/mapping_options.h"

  // 1. Define and setup the TC compilation unit with CUDA memory management backed by ATen.
  std::string tc = R"TC(
  def TensorDot(float(N, C1, C2, H, W) I0, float(N, C2, C3, H, W) I1) -> (O) {
    O(n, c1, c3, h, w) +=! I0(n, c1, c2, h, w) * I1(n, c2, c3, h, w)
  })TC";

  // 2. Allocate tensors with random data
  at::Tensor I0 = at::CUDA(at::kFloat).rand({32, 512, 8, 28, 28});
  at::Tensor I1 = at::CUDA(at::kFloat).rand({32,   8, 2, 28, 28});
  std::vector<at::Tensor> outputs;

  // 3. Run autotuning with evolutionary search starting from a naive option
  auto options = tc::MappingOptions::makeNaiveMappingOptions();
  auto bestOption = autotune(cacheFilename, tc, "TensorDot", {I0, I1}, options, {options});

  // 4. Compile and run the TC with the best option.
  tc::ATenCompilationUnit atCompl;
  atCompl.define(tc);
  auto handle = atCompl.compile("TensorDot", {I0, I1}, bestOption);
  atCompl.run("TensorDot", {I0, I1}, outputs, handle);

  // 5. Perform precision checks against an ATen reference implementation
  check({I0, I1}, outputs, [&I0, &I1](){ return ...; });