基于WebGL实现矩阵计算

Posted 2021-04-25 神机喵算

节选自《深度学习TensorFlow.js：浏览器实战篇》第五章，已获授权。

基于纹理和着色器的矩阵计算

学完WebGL的GPGPU之后，我们回到深度学习，构建一个简单的WebGL线性代数库WGLMatrix。然后我们使用WGLMatrix实现一个简单的神经网络（不是卷积神经网络）从MNIST数据集中学习识别手写数字。线性代数库的第一个测试版在代码仓库chapter5/4_WGLMatrix。

标准的矩阵加法

GLSL中的内建矩阵类型的最大维度不多于四维，这对深度学习是不够的。所以用纹理存储矩阵，我们需要开发自定义的着色器程序进行矩阵操作。每个纹理元素时一个矩阵项，纹理的分辨率和矩阵的维度一样大小。

当创建一个矩阵，如果用四个javascript初始化的数组，RGBA通道将用这些来填充。如果只有一个数组，那它的值将会重复四次以填充颜色通道。每个常用的矩阵计算操作对于RGBA通道来说，是分开处理的。就好像用四个不同的矩阵并行的处理四次。

矩阵加法的片段着色器程序像其它对位的操作一样，不依赖矩阵的大小，下面是对应的GLSL代码：

void main(void){ vec2 uv=gl_FragCoord.xy/resolution; vec4 matAValue=texture2D(samplerTexture0, uv); vec4 matBValue=texture2D(samplerTexture1, uv); gl_FragColor=matAValue+matBValue;} 

标准的矩阵乘法：

和矩阵加法的着色器程序不同的是，矩阵乘法的着色器用for循环将左矩阵第一行乘以右矩阵第一列。在WebGL1中for条件不能使用非常数值。该限制包括GLSL的 uniform类型值或者上一步计算的值。所以，我们需要为每个矩阵乘法编译一个着色器程序。例如，下面的着色器程序可以进行所有的（n, 10）矩阵和（10，n）矩阵相乘：

//vector between 2 consecutive texels of first factor:const vec2 DU=vec2(1./10., 0.);//vector between 2 consecutive texels of second factor: const vec2 DV=vec2(0., 1./10.);
void main(void){vec2 uv=gl_FragCoord.xy/resolution; vec2 uvu=uv*vec2(1.,0.);vec2 uvv=uv*vec2(0.,1.);vec4 result=vec4(0.,0.,0.,0.);for (float i=0.0; i<10.0; i+=1.0){ result+=texture2D(samplerTexture0, uvv+(i+0.5)*DU) *texture2D(samplerTexture1, uvu+(i+0.5)*DV);} gl_FragColor=result; }

上面的代码中，对于每个i增加0.5是为了获取像素的中间值，反之，对于某些矩阵维度可能会出现舍入错误。在WebGL2中，for循环条件实现了非常数条件。但是，这对于只支持WebGL2的商业应用是不能接受的（在2018年3月，支持率是41%，数据来源于webglstats.com）。所以我们构建的WebGL1着色器程序也支持WebGL2.

我们经常会同时做矩阵相乘和矩阵相加。例如，我们将权重矩阵和神经网络layer输入矩阵X相乘，然后增加偏置矩阵B求和得到输入矩阵Z：Z ＝ WX ＋ B。我们应该编写一个自定义的着色器程序实现该功能，称为FMA（Fused Multiply–Accumulate）。它将会节省一些渲染到纹理的过程。

矩阵库WGLMatrix管理矩阵乘法和FMA着色器程序的一个字典。如果我们需要处理两个矩阵，其中一个矩阵的维度包含在字典，则我们只需要用字典中的着色器程序。否则，需要重新编写新的维度着色器程序，并增加到字典。

激活函数

着色器程序致力于激活函数的应用。我们需要注意浮点型特殊值，特别地，激活函数涉及到指数函数或者对数函数。下面的着色器程序应用于sigmoid激活函数：

const vec4 ONE=vec4(1.,1.,1.,1.); void main(void) {vec2 uv=gl_FragCoord.xy/resolution; vec4 x=texture2D(samplerTexture0, uv); vec4 y;y=1./(ONE+exp(-x)); gl_FragColor=y; }

因为许多自定义的激活函数都被应用于矩阵，所以我们在WGLMatrix库设置公共方法来编写自定义着色器程序。

运用WGLMatrix库

每个矩阵在使用之前都会初始化。用扁平的值分别去初始化矩阵m、v、n：

// encoding 3*3 matrix | 0 1 2 |:// | 3 4 5 |// | 6 7 8 | varM=newWGLMatrix.Matrix(3,3,[0,1,2, 3,4,5, 6,7,8]); //V is a column matrix, which is a vector:var V=new WGLMatrix.Matrix(3,1,[1,2,3]); var W=new WGLMatrix.MatrixZero(3,1);

从数学上讲，一个向量和一个矩阵没有什么区别，维度为1。初始化之后我们可以进行矩阵操作。为了对矩阵A执行OPERATION操作，执行如下代码：

A.OPERATION(arguments..., R)

R是计算结果矩阵。我们不能将操作结果存储在所涉及的任何矩阵中，因为不可能同时读取一个纹理并渲染到该纹理。所以操作经常返回结果矩阵R。

例如，为了进行操作W＝M＊V，执行下面的代码：

M.multiply(V, W);

上面的代码返回矩阵W。我们可以声明一个自定义的对位操作，执行如下代码：

WGLMatrix.addFunction('y=cos(x);', 'COS');

上面代码中第一个参数是函数的GLSL代码，使用预定义的向量vec4 x和vec4 y。第二个参数是函数的标记。然后将其应用于矩阵M的所有元素：M.apply('COS', R)。其中R是返回的结果矩阵。

---

若发现以上文章有任何不妥，请联系我。