python - 如何矢量化 torch 函数？

Question

当使用numpy时，我可以使用np.vectorize来向量化包含if语句的函数，以便该函数接受数组参数。如何使用 torch 执行相同操作，以便函数接受张量参数？

例如，下面代码中的最终打印语句将失败。我怎样才能做到这一点？

import numpy as np
import torch as tc

def numpy_func(x):
    return x if x > 0. else 0.
numpy_func = np.vectorize(numpy_func)

print('numpy function (scalar):', numpy_func(-1.))
print('numpy function (array):', numpy_func(np.array([-1., 0., 1.])))

def torch_func(x):
    return x if x > 0. else 0.

print('torch function (scalar):', torch_func(-1.))
print('torch function (tensor):', torch_func(tc.tensor([-1., 0., 1.])))

Answer 1

您可以使用 .apply_() 作为 CPU 张量。对于 CUDA 来说，这项任务是有问题的:if 语句不容易进行 SIMDify。

您可以对 functorch.vmap 应用与用于着色器的视频驱动程序相同的解决方法:评估条件的两个分支并坚持算术。

否则，只需使用 for 循环:这就是 np.vectorize() 的主要用途 does anyway .

def torch_vectorize(f, inplace=False):
    def wrapper(tensor):
        out = tensor if inplace else tensor.clone()
        view = out.flatten()
        for i, x in enumerate(view):
            view[i] = f(x)
        return out
    return wrapper