Facebook 开源数学语言 Tensor Comprehensions Facebook 开源数学语言 Tensor Compreh

Facebook AI 研究院近日开源了 Tensor Comprehensions ，这是一个 C++ 库及数学语言，功能齐全，能有效填补研究人员于数学运算领域的沟通鸿沟，并基于各种硬件后端上大规模运行工程模型。

Tensor Comprehensions 采用了 Just-In-Time 的编译自动生成机器学习社区所需的高性能代码，并被设计为高度可移植的。通过 Tensor Comprehensions，研究人员能够以数学符号的方式进行编写，系统能够根据需求进行编译调整，并输出专业的代码。

示例：

#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 ...; });

大致结果如下：