原语库

深度学习框架是如何提供灵活性的？

import torch
from torch.autograd import Variable

x = torch.randn((5), requires_grad=True)

mean = torch.randn((5))
var = torch.randn((5))

with torch.autograd.profiler.profile(enabled=True, use_cuda=True) as prof:
    y = (x-mean)/var
    y.backward(torch.ones_like(y))
prof.export_chrome_trace("profile.json")

STAGE:2023-07-10 13:28:17 3224314:3224314 ActivityProfilerController.cpp:311] Completed Stage: Warm Up
STAGE:2023-07-10 13:28:26 3224314:3224314 ActivityProfilerController.cpp:317] Completed Stage: Collection
STAGE:2023-07-10 13:28:26 3224314:3224314 ActivityProfilerController.cpp:321] Completed Stage: Post Processing

这段计算是没有计算库与之对应的，深度学习框架一般的做法是先让芯片供应商提供 primitive 库，包含基本的四则运算、卷积、矩阵乘法和访存操作（事实上，primitive 越完备，可以支持的计算越多）。比如：

def add(a, b):
def sub(a, b):
def mul(a, b):
def div(a, b):

有了 primitive 库，当遇到一段计算的时候，会将计算分解由 primitive 来表达。这样，尽管我们并没有直接支持这个计算，但是它可以通过我们调用 primitive 来实现。到这里，灵活性的问题基本解决了。如果没解决，那就多加 primitive，同时尝试支持 loop、recursive、branch 和 closure 等。