使用taichi 写了一个任意平台任意显卡推理的Linear

这东西就是在于任意的显卡都能加速任意模型
当然如何有人使用taichi写一个卷积那么计算机视觉也能任意显卡加速人工智能
如果还有人写了个深度学习训练框架那么恭喜AMD,ARM 等任何芯片厂商乐疯

import taichi as ti
import numpy as np
ti.init(arch=ti.vulkan)


class Linear():
    def __init__(self, input_size, output_size, weights=None):
        if weights:
            self.weights = weights
        else:
            # self.weights = ti.Matrix([[0] * output_size] * input_size)
            self.weights = ti.Matrix(np.random.random(output_size*input_size).reshape([input_size,output_size]))

    @staticmethod
    def taichi_mul(a, b):
        ar, al = a.to_numpy().shape
        br, bl = b.to_numpy().shape
        assert al == br

        @ti.kernel
        def mlp() -> ti.types.matrix(ar, bl, dtype=ti.float32):
            return a @ b

        return mlp()

    def forward(self, a):
        return self.taichi_mul(a, self.weights)

    def set_weights(self, weights):
        self.weights = ti.Matrix(weights)


l1 = Linear(3, 5)
print(l1.forward(ti.Matrix([[1] * 3])))
这段代码实现了一个简单的线性层(Linear layer)类,通过调用 forward 函数可以对输入进行线性变换。其中 Linear 类中的 weights 表示线性层的权重矩阵,可以通过构造函数的输入或者 set_weights 函数进行设置。forward 函数使用 taichi_mul 函数实现了矩阵乘法,并返回了乘积结果。

在 __init__ 函数中,如果 weights 参数被传入,则将其作为权重矩阵;否则通过 np.random.random 函数生成随机的权重矩阵。其中,输入参数 input_size 和 output_size 分别表示输入和输出的特征数。

在 taichi_mul 函数中,输入参数 a 和 b 分别表示两个矩阵,通过 Taichi 提供的 @ 运算符实现了矩阵乘法,返回乘积结果。值得注意的是,Taichi 不能直接操作 Python 中的数据类型,因此在使用 Taichi 前,需要将 Python 中的数据类型转换为 Taichi 中的数据类型,可以调用 to_numpy() 函数将 Taichi 中的数据类型转换为 NumPy 中的数据类型,然后对其进行操作,最后再调用 ti.Matrix() 函数将其转换为 Taichi 中的数据类型。

在主函数中,首先构造了一个输入矩阵 [[1] * 3],然后通过 l1.forward() 函数将其输入到 l1 线性层中,得到线性变换的结果,最后将结果打印输出。

你可能感兴趣的:(taichi,windows,重构,transformer,神经网络,前端)