Taichi能够更精细地控制并行和每个元素(element)的操作,极大地提升了用户操作的灵活性。而Torch则将这些细节抽象成张量(Tensor)级别的操作,使得用户能聚焦于机器学习的模型结构。
01
作为机器学习、计算机图形学领域炙手可热的框架和编程语言,Torch 和 Taichi 能否各取所长,结合使用呢? 答案是肯定的。在本篇文章中,作者将通过两个简单的例子演示:如何使用 Taichi Kernel 来实现 PyTorch 程序中特殊的数据预处理和自定义的算子,告别手写 CUDA,用轻巧便捷的方式提升机器学习模型算法的开发效率和灵活性。
案例 1:数据预处理
边缘填充(Padding)是机器学习中常用的预处理方法。如在对图像执行卷积操作时,用户需要对图像边缘进行填充,以保证图像输入输出前后的尺寸不变。一般来说,填充的方法有零填充或 torch.nn.functional.pad 提供的重复填充、循环填充等其他预设模式。但有时候我们想要在边缘上填充某个特殊的纹理或者模式,却并没有一个精心优化过的 PyTorch 算子能够适配这种场景。
解决方案有两个:使用 PyTorch 或者 Python 逐个操作矩阵元素;手写 C++ 或 CUDA 代码并接入PyTorch。前者的计算效率非常低,会拖累神经网络的训练速度;后者学习曲线陡峭,实操非常麻烦,开发流程冗长。
那么,有没有更好的方案呢?接下来我们将通过一个例子,带大家体验如何用 Taichi 做一个砖墙纹理的边缘填充。
用Taichi给PyTorch「添砖加瓦」!
第一步,我们在PyTorch中创建一个如下图所示的「砖块」。为了更好地观察填充的规律,我们给这块「砖」填充上了渐变的颜色:
填充的基本单元
第二步,我们想要在x轴上错位重复这个「砖」,也就是如下所示的效果:
由于PyTorch中没有为这样的填充提供原生的算子,为了提高运算效率,需要将padding过程改写成一系列PyTorch的原生矩阵运算:
def torch_pad(arr, tile, y ):
# image_pixel_to_coord
arr[:, :, 0 ] = image_height - 1 + ph - arr[:, :, 0 ]
arr[:, :, 1 ] -= pw
arr1 = torch.flip(arr, ( 2 , ))
# map_coord
v = torch. f