创建自定义 WebGL Ops（操作）
Posted 黑胡桃实验室
tags:
篇首语：本文由小常识网(cha138.com)小编为大家整理，主要介绍了创建自定义 WebGL Ops（操作）相关的知识，希望对你有一定的参考价值。
要创建自定义 WebGL Ops（操作），我们要做的是创建一个对象实现接口 tf.webgl.GPGPUProgram。
该接口定义为：
interface GPGPUProgram {
variableNames: string[];
outputShape: number[];
userCode: string;
supportsBroadcasting?: boolean;
}
举一个普通的例子，让我们实现一个 Ops（操作）f(x) = x * x + x。
GLSL 代码如下：
void main() {
float x = getXAtOutCoords();
float value = x * x + x;
setOutput(value);
}
其中 getXAtOutCoords 获取输入值，setOutput 返回结果值。需要注意的是，输出 Tensor（张量）中的每一个值都将调用 main 函数。
完整的 GPGPUProgram 定义是：
const squareAndAddKernel = inputShape => ({
variableNames: ['X'],
outputShape: inputShape.slice(),
userCode: `
void main() {
float x = getXAtOutCoords();
float value = x * x + x;
setOutput(value);
}
`
})
要运行此操作，您可以使用 tf.ENV.backend.compileAndRun(program: GPGPUProgram, inputs: tf.Tensor[]): tf.Tensor。请注意，后端应为 WebGL，否则将会被判断为“Undefined”。
const x = tf.tensor([1, 2, 3, 4]);
const program = squareAndAddKernel(x.shape);
const result = tf.ENV.backend.compileAndRun(program, [x]);
除此之外，我们可能还想为此操作定义 gradient（梯度），以便 gradient（梯度）可以通过它反向传播。
为此，我们使用 tf.customGrad。
const squareAndAddBackpropKernel = inputShape => ({
variableNames: ['X'],
outputShape: inputShape.slice(),
userCode: `
void main() {
float x = getXAtOutCoords();
float value = 2.0 * x + 1.0;
setOutput(value);
}
`
});
const squareAndAdd = tf.customGrad(x => {
const backend = tf.ENV.backend;
const program = squareAndAddKernel(x.shape);
const backpropProgram = squareAndAddBackpropKernel(x.shape);
const value = backend.compileAndRun(program, [x]);
const gradFunc = dy =>
[backend.compileAndRun(backpropProgram, [x]).mul(dy)];
return {value, gradFunc};
});
然后我们可以将其用于：
const x = tf.tensor([1, 2, 3, 4]);
const value = squareAndAdd(x);
const grads = tf.grad(x => squareAndAdd(x));
const dx = grads(x);
// value == [2, 6, 12, 20]
// dx == [3, 5, 7, 9]
或者更简洁的：
const {value, grad} = tf.valueAndGrad(squareAndAdd)(x);
由 TensorFlow.js 生成的 GLSL 函数
TensorFlow.js 能够生成可用于读取输入 Tensor（张量）并写入输出 Tensor（张量）的函数，以及其他数值操作的实用函数。这些都会由 着色编译器 预先添加到您的代码中。
1void setOutput(float value)
设置运行片段着色器的坐标的输出值（相当于 gl_FragCoord = vec4(value, 0.0, 0.0, 0.0)）。
2indexType getOutputCoords()
其中 indexType 是 int | ivec2 | ivec3 | ivec4 | ivec5 | ivec6 中的一个。
如果输出 Tensor 为 rank-0 或 rank-1，则返回值为 int，否则返回 ivecN，其中 N == rank 。这是该程序将要写入的输出 Tensor（张量）中的单元格坐标。
3Tensorflow.js 生成 GLSL 函数通过输入 Tensor（张量）进行采样
形式如下：
float get{VarName}AtOutCoords()
float get{VarName}() // 输入为 rank-0
float get{VarName}(int x) // 输入为 rank-1
float get{VarName}(int x, int y) // 输入为 rank-2
float get{VarName}(int x, int y, int z) // 输入为 rank-3
float get{VarName}(int x, int y, int z, int w) // 输入为 rank-4
// rank-5 型输入和 rank-6 型输入继续上述步骤
// 例如, 对于名为“x”的 rank-2 的 Tensor（张量）：
// float getX(int x, int y)
其中 VarName 是在 GPGPUProgram 的 variableNames 数组中定义的变量名，首字母大写。这意味着对于名为 matrix 的变量，TensorFlow.js 将生成 getMatrix。
其中许多函数都取决于输入 Tensor（张量）的 rank (秩)，所以在您的 GPGPUProgram 中，您应该根据 inputShape 的 rank (秩)提供不同的代码。例如，如果 get{VarName}AtOutCoords() 不存在，我们可能将 squareAndAddKernel 编写为：
const squareAndAddKernel = inputShape => ({
const variableNames = ['X']
const outputShape = inputShape.slice()
const rank = outputShape.length
const coordSnippets = ['',
'coords',
'coords.x, coords.y',
'coords.x, coords.y, coords.z',
'coords.x, coords.y, coords.z, coords.w']
const coordType = rank < 2 ? 'int' : `ivec${rank}`
const userCode = `
void main() {
${coordType} coords = getOutputCoords();
float x = getX(${coordSnippets[rank]});
setOutput(x * x + x);
}`
return {variableNames, outputShape, userCode}
})
bool isNaN(float val)：如果 val 为 NaN，返回 true，否则返回 false。
int round(float value)：将 value 化整为最接近它自己的整数。
int imod(int x, int y)：与 float mod(float x, float y) 相同，但除了 ints，因为 GLSL 没有提供给我们。
float random(float seed)：根据 https://www.shadertoy.com/view/4djSRW 中的 Dave Hoskins 公式返回一个伪随机数。
✄----------------------------------
本文是黑胡桃实验室与 Googler、GDE 协作翻译的 TensorFlow.js 中文尝鲜版，英文版首发于 https://js.tensorflow.org/
探索AI
触摸科技
黑胡桃实验室
杭州·赛银国际广场4幢1楼
BlackWalnut Labs.
以上是关于创建自定义 WebGL Ops（操作）的主要内容，如果未能解决你的问题，请参考以下文章
VSCode自定义代码片段15——git命令操作一个完整流程