前言
TensorFlow.js 是一个用于使用 JavaScript 进行机器学习开发的库。TensorFlow.js 支持可实现张量存储和数学运算的多种不同后端。在任何给定时间内,均只有一个后端处于活动状态。在大多数情况下,TensorFlow.js 会根据当前环境自动为您选择最佳后端。WebGL 是当前适用于浏览器功能最强大的后端,此后端的速度比普通 CPU 后端快 100 倍。
除了 TensorFlow.js,还有 GPU.js 、Babylon.js ****等库也采用 WebGL 做计算加速。为什么 WebGL 可以获得如此高的计算速度?
GPU 计算
在解释 WebGL 为什么可以被用来做计算之前,需要先回答为什么 GPU 被用来做计算?
在我们常规的认知中,CPU 是被用来做计算,而 GPU 是用来做图形渲染的。得益于现代图形处理器有强大的并行处理能力和可编程流水线,令图形处理器也可以处理非图形数据。特别是在面对单指令流多数据流(SIMD)且数据处理的运算量远大于数据调度和传输的需要时,通用图形处理器在性能上大大超越了传统的中央处理器应用程序。
举个例子,矩阵相乘运算:
其中:
因此,可以把矩阵运算,拆分成 n*n 个计算任务,当每个计算任务都完成后,综合在一起就是一个完成的结果。用通俗的话说,如果一个计算任务可以被拆分成许多个重复计算,那么 GPU 非常适合干,而 CPU 适合做复杂的计算任务。
利用处理图形任务的图形处理器来计算原本由中央处理器处理的通用计算任务,这被称为图形处理单元上的通用计算,简称 GPGPU(General-purpose computing on graphics processing units)。
WebGL 基础知识
WebGL 通常被认为是一种3D API,通过 WebGL 可以绘制 3D 图形,但实际上 WebGL 仅仅是一个栅格化引擎,基于代码绘制点、线、三角形,用户基于点、线、三角形组合完成 3D 任务,所以它是一个很底层的图形 API。
工作原理
WebGL 最终运行在 GPU 上,代码用 GLSL 语言(图形库着色器语言)编写,主要由两部分组成,分别是顶点着色器(Vetex Shader)和片段着色器(Fragment Shader),两者组合起来成为一个程序(Program)。
- 顶点着色器:计算点的位置。
- 片段着色器:计算当前正在绘制图形的每个像素的颜色。
几乎所有的 WebGL API 是为这两个着色器运行来设置状态,然后调用 gl.drawArrays 和 gl.drawElements 在 GPU上运行着色器。
根据官方文档的示例,假设要绘制下面的渐变三角形,WebGL 的工作流程:
1.先通过顶点着色器确定其中一个顶点的位置
2.顶点位置传给片段着色器,确定该顶点的颜色
3.通过顶点着色器确定第二个顶点的位置
4.顶点位置传给片段着色器,确定该顶点的颜色
5.通过顶点着色器确定第三个顶点的位置
6.顶点位置传给片段着色器,确定该顶点的颜色
7.栅格化整个三角形,基于三个顶点的颜色值做插值计算
各像素点栅格化并行执行。
想了解更多关于 WebGL 的工作原理,可以参考:WebGL2的基本原理、WebGL2 是如何工作的。
WebglProgram 基类
WebGL 的工作流程是定死的,可以变化的是着色器代码,因此可以抽离出一个 WebGL 的程序基类,其包含的功能:
- init 方法:根据顶点着色器和片段着色器创建着色程序
- render 方法:进行图形渲染
- read 方法:获取当前 canvas 渲染的图像中的像素数据
class WebglProgram {
constructor(dimen, canvasSize) {
this.dimen = dimen;
this.canvas = document.createElement("canvas");
this.canvas.width = canvasSize || dimen;
this.canvas.height = canvasSize || dimen;
this.gl = this.canvas.getContext("webgl2");
this.program = this.gl.createProgram();
}
/**
* 程序初始化
* @param {*} vertexShader
* @param {*} fragmentShader
*/
async init(vertexShader, fragmentShader) {
const vshaderCode = await this.loadRes(vertexShader);
let fshaderCode = await this.loadRes(fragmentShader);
fshaderCode = fshaderCode.replace(/CANVAS_SIZE/g, this.dimen);
this.initShader(vshaderCode, this.gl.VERTEX_SHADER);
this.initShader(fshaderCode, this.gl.FRAGMENT_SHADER);
this.gl.linkProgram(this.program);
this.gl.useProgram(this.program);
this.gl.bindBuffer(this.gl.ARRAY_BUFFER, this.gl.createBuffer());
let vecPosXArr = new Float32Array([-1, -1, 1, -1, -1, 1, 1, 1]);
this.gl.bufferData(this.gl.ARRAY_BUFFER, vecPosXArr, this.gl.STATIC_DRAW);
let posAtrLoc = this.getAttribLoc("g_pos");
this.gl.enableVertexAttribArray(posAtrLoc);
this.gl.vertexAttribPointer(posAtrLoc, 2, this.gl.FLOAT, false, 0, 0);
this.gl.clearColor(0.0, 0.0, 0.0, 1);
this.gl.clear(this.gl.COLOR_BUFFER_BIT);
}
/**
* 创建着色器
* @param {着色器代码} code
* @param {着色器类型} type
*/
initShader(code, type) {
const shader = this.gl.createShader(type);
this.gl.shaderSource(shader, code);
this.gl.compileShader(shader);
if (!this.gl.getShaderParameter(shader, this.gl.COMPILE_STATUS))
throw new Error("compile: " + this.gl.getShaderInfoLog(shader));
this.gl.attachShader(this.program, shader);
}
/**
* 创建纹理
* @param {纹理索引} index
* @param {*} tSampler
* @param {纹理数据} pixels
*/
initTexture(index, tSampler, pixels) {
const texture = this.gl.createTexture();
this.gl.activeTexture(this.gl[`TEXTURE${index}`]);
this.gl.bindTexture(this.gl.TEXTURE_2D, texture);
this.gl.texParameteri(
this.gl.TEXTURE_2D,
this.gl.TEXTURE_MIN_FILTER,
this.gl.LINEAR
);
this.gl.texParameteri(
this.gl.TEXTURE_2D,
this.gl.TEXTURE_MAG_FILTER,
this.gl.NEAREST
);
this.gl.texImage2D(
this.gl.TEXTURE_2D,
0,
this.gl.RGBA,
this.dimen,
this.dimen,
0,
this.gl.RGBA,
this.gl.UNSIGNED_BYTE,
pixels,
0
);
// 使得着色器中通过 ‘tSampler’ 访问到数据
this.gl.uniform1i(this.getUniformLoc(tSampler), index);
}
/**
* 给 Uniform 变量赋值
* @param {*} tUniform
* @param {*} value
*/
initUniform(tUniform, value) {
const uniLoc = this.getUniformLoc(tUniform);
this.gl.uniform1fv(uniLoc, value);
}
/**
* 获取属性变量
* @param {*} name
* @returns
*/
getAttribLoc(name) {
let loc = this.gl.getAttribLocation(this.program, name);
if (loc == -1) throw `getAttribLoc ${name} error`;
return loc;
}
/**
* 获取 Uniform 变量
* @param {*} name
* @returns
*/
getUniformLoc(name) {
let loc = this.gl.getUniformLocation(this.program, name);
if (loc == null) throw `getUniformLoc ${name} err`;
return loc;
}
/**
* 获取着色器文本
* @param {*} file
* @returns
*/
async loadRes(file) {
const resp = await fetch(file);
return resp.text();
}
/**
* 图形绘制
*/
render() {
this.gl.drawArrays(this.gl.TRIANGLE_STRIP, 0, 4);
}
/**
* canvas 读取像素数据
* @returns
*/
read() {
let picBuf = new ArrayBuffer(this.canvas.width * this.canvas.height * 4);
let picU8 = new Uint8Array(picBuf);
let picU32 = new Uint32Array(picBuf);
this.gl.readPixels(
0,
0,
this.canvas.width,
this.canvas.height,
this.gl.RGBA,
this.gl.UNSIGNED_BYTE,
picU8
);
return picU32;
}
}
纹理基础
在 WebGL 中,着色器有四种获取数据的方法,分别是:
- 属性(Attributes),缓冲区(Buffers)和顶点数组(Vetex Arrays)
- 全局变量(Uniforms)
- 纹理(Textures)
- 可变量(Varyings)
纹理是能够在着色器程序中随机访问的数组数据。大多数情况下纹理存储图片数据,但它也用于包含颜色以外的数据。
创建纹理的基本流程大致有以下几步:
// 1.激活纹理单元
gl.activeTexture(...);
// 2.创建和绑定
const texture = gl.createTexture();
gl.bindTexture(gl.TEXTURE_2D, texture);
// 3.参数设置
gl.texParameteri(...);
gl.texParameteri(...);
// 4.填充纹理的内容
gl.texImage2D(..., image);
// 5.通过纹理单元将纹理传送给着色器程序
gl.uniform1i(...);
实际使用可以参照上面基类 WebglProgram 中的initTexture方法。
纹理使用案例
下面看一个纹理使用案例:
- 1.先创建一个顶点着色器:
此处顶点着色器的作用是进行坐标转换,在着色器中坐标范围为【-1,1】,转换后变成【0,1】,其中左下角为【0,0】
#version 300 es
precision highp float;
precision highp int;
in vec4 g_pos;
out vec2 v_pos;
void main() {
float curX = (g_pos.x + 1.) / 2.;
float curY = (g_pos.y + 1.) / 2.;
v_pos = vec2(curX, curY);
gl_Position = g_pos;
}
- 2.再创建一个片段着色器:
WebGL 在处理图片纹理编写片段着色器的时候,会通过 sampler2D 关键字声明一个纹理相关的变量,sampler2D 和 vec2、float 表示一种数据类型,只不过 sampler2D 关键字声明的变量表示一种取样器类型变量,简单点说就是该变量对应纹理图片的像素数据。
WebGL 着色器提供了内置函数 texture2D,可以直接使用,通过 texture2D/texture 函数,可以从纹理图像提取像素值。
#version 300 es
precision highp float;
precision highp int;
in vec2 v_pos;
uniform sampler2D samplerA;
out vec4 o_result;
vec4 reverse(const vec4 v){
return v.abgr;
}
void main() {
vec4 color = texture(samplerA, v_pos);
o_result = reverse(color);
}
- 3.设置初始纹理数据,渲染图形
/**
* 创建矩阵,以一维数组代替
* @param {矩阵维数} dims
* @param {矩阵中每一项生成的规则} fn
* @returns
*/
function createMatrix(dims, fn) {
let matrix = new Uint32Array(dims * dims);
for (let i = 0; i < dims; i++) {
for (let j = 0; j < dims; j++) {
matrix[i * dims + j] = fn(i, j);
}
}
return matrix;
}
class TextureDemo extends WebglProgram {
constructor(dimen, canvasSize) {
super(dimen, canvasSize);
document.body.appendChild(this.canvas);
}
async init(matrixA) {
await super.init("demo2_v_shader.c", "demo2_f_shader.c");
this.initTexture(0, "samplerA", matrixA);
}
}
async function ShowTextureDemo() {
const colorMap = [
// 红、绿、蓝
[0xff0000ff, 0x00ff00ff, 0x0000ffff],
[0xffff00ff, 0xff00ffff, 0x00ffffff],
[0x000000ff, 0xffffffff, 0xf0f0f0ff],
];
const randomMatrix = createMatrix(colorMap.length, (x, y) => colorMap[x][y]);
const randomMatrixU8 = new Uint8Array(randomMatrix.buffer);
const demoShowTexture = new TextureDemo(colorMap.length, 200);
// changeBufferToGRBA(randomMatrixU8);
await demoShowTexture.init(randomMatrixU8);
demoShowTexture.render();
}
ShowTextureDemo()
- 4.渲染结果
特别注意:在片段着色器代码中有对 color 做一次顺序反转,从 rgba 转换成 abgr ,原因是在从 Uint32Array 转成 Uint8Array 时会导致数据顺序反转,原因参见:Javascript Typed Arrays and Endianness。实测效果如下:
WebGL 计算案例
WebGL 通过 canvas 画布展示图形,而 canvas 由像素点组成,每个像素点有 RGBA 四个维度,对应的数字是0-255,即每个维度 8 位,所以从二进制角度考虑,一个像素正好 32 位,可以被用来作为一个 int 类型,这是 WebGL 被用来做计算的核心---用像素存储数据。
此处,以上面的矩阵相乘作为计算案例。
CPU 进行矩阵计算
/**
* 创建矩阵,以一维数组代替
* @param {矩阵维数} dims
* @param {矩阵中每一项生成的规则} fn
* @returns
*/
function createMatrix(dims, fn) {
let matrix = new Uint32Array(dims * dims);
for (let i = 0; i < dims; i++) {
for (let j = 0; j < dims; j++) {
matrix[i * dims + j] = fn(i, j);
}
}
return matrix;
}
// 矩阵相乘
const matrixMultiply = function(dimensions,ma, mb) {
return createMatrix(dimensions, function(x, y) {
let sum = 0;
for (let i = 0; i < dimensions; i++) {
sum += ma[x * dimensions + i] * mb[i * dimensions + y];
}
return sum;
});
}
function matrixMultiplyCPU(dimensions) {
// 随机生成 n 维的矩阵
const randomMatrix = createMatrix(dimensions, () => Math.floor(Math.random() * 1000));
// 计算矩阵的平方
console.time("CPU 计算");
result = matrixMultiply(dimensions,randomMatrix, randomMatrix);
console.timeEnd("CPU 计算");
}
matrixMultiplyCPU(300) // matrixMultiplyCPU: 62.09ms
WebGL 进行矩阵计算
- 1.创建顶点着色器:
#version 300 es
precision highp float;
precision highp int;
in vec4 g_pos;
out vec2 v_pos;
void main() {
float curX = (g_pos.x + 1.) / 2.;
float curY = (g_pos.y + 1.) / 2.;
v_pos = vec2(curX, curY);
gl_Position = g_pos;
}
- 2.创建片段着色器
#version 300 es
precision highp float;
precision highp int;
const int U_LENGTH = CANVAS_SIZE;
const float U_TEXTURE_POS_FIX = .5 / float(U_LENGTH);
in vec2 v_pos;
uniform sampler2D samplerA;
uniform sampler2D samplerB;
out vec4 o_result;
vec4 int2rgba(const int i) {
vec4 v4;
v4.r = float(i >> 24 & 0xFF) / 255.;
v4.g = float(i >> 16 & 0xFF) / 255.;
v4.b = float(i >> 8 & 0xFF) / 255.;
v4.a = float(i >> 0 & 0xFF) / 255.;
return v4;
}
int rgba2int(const vec4 v) {
int r = int(v.r * 255.) << 24;
int g = int(v.g * 255.) << 16;
int b = int(v.b * 255.) << 8;
int a = int(v.a * 255.) << 0;
return r + g + b + a;
}
vec4 reverse(const vec4 v){
return v.abgr;
}
int getMaxtrixValue(const sampler2D sampler, const float x, const float y){
vec4 pixel = texture(sampler, vec2(x, y));
return rgba2int(reverse(pixel));
}
void main() {
int sum = 0;
float textPos = 0.0;
for (int i = 0; i < U_LENGTH; i++) {
textPos = U_TEXTURE_POS_FIX + float(i) / float(U_LENGTH);
sum += getMaxtrixValue(samplerA, v_pos.x, textPos) * getMaxtrixValue(samplerB, textPos, v_pos.y);
}
o_result = reverse(int2rgba(sum));
}
- 3.设置纹理数据,渲染图形:
class MatrixTexture extends WebglProgram {
async init(matrixA, matrixB) {
await super.init("matrix_v_shader.c", "matrix_f_shader.c");
this.initTexture(0, "samplerA", matrixA);
this.initTexture(1, "samplerB", matrixB);
}
}
async function showMatrixCompute(dimensions) {
console.log("dimensions", dimensions);
const randomMatrix = createMatrix(dimensions, () => Math.floor(Math.random() * 1000));
const randomMatrixU8 = new Uint8Array(randomMatrix.buffer);
console.time("WebGL 计算");
const showMatrixTexture = new MatrixTexture(dimensions);
await showMatrixTexture.init(randomMatrixU8, randomMatrixU8);
showMatrixTexture.render();
const result = showMatrixTexture.read();
console.timeEnd("WebGL 计算");
}
showMatrixCompute(100)
- 4.计算结果对比:
举个 🌰 说明
- 1.随机创建一个 2 *2 矩阵,类型是 Uint32Array
const randomMatrix = new Uint32Array([908, 766, 271, 434]);
// Uint32Array(4) [908, 766, 271, 434, buffer: ArrayBuffer(16), byteLength: 16, byteOffset: 0, length: 4, Symbol(Symbol.toStringTag): 'Uint32Array']
- 2.像素点由 4 个 Uint8Array 组成,所以转成 Uint8Array,经历一次 reverse
const randomMatrixU8 = new Uint8Array(randomMatrix.buffer);
// Uint8Array(16) [140, 3, 0, 0, 254, 2, 0, 0, 15, 1, 0, 0, 178, 1, 0, 0, buffer: ArrayBuffer(16), byteLength: 16, byteOffset: 0, length: 16, Symbol(Symbol.toStringTag): 'Uint8Array']
/* reverse 说明
140 --转成二进制-> 10001100
3 --转成二进制-> 00000011
00000011 10001100 --转成十进制-> 908
*/
- 3.将上面的数据作为纹理数据传入 WebGL 程序
此时的纹理数据:
[140, 3, 0, 0, 254, 2, 0, 0, 15, 1, 0, 0, 178, 1, 0, 0]
PS: 实际在转换时要做的处理:归一化、reverse
如果把纹理绘制出来,效果如下:
- 4.WebGL 程序从纹理采样器中获取数据并计算
int getMaxtrixValue(const sampler2D sampler, const float x, const float y){
vec4 pixel = texture(sampler, vec2(x, y));
return rgba2int(reverse(pixel));
}
void main() {
int sum = 0;
float textPos = 0.0;
for (int i = 0; i < U_LENGTH; i++) {
textPos = U_TEXTURE_POS_FIX + float(i) / float(U_LENGTH);
sum += getMaxtrixValue(samplerA, v_pos.x, textPos) * getMaxtrixValue(samplerB, textPos, v_pos.y);
}
o_result = reverse(int2rgba(sum));
}
以左下角像素值计算为例,详细计算过程如下,:
-
- 先从纹理A中取左下角(0,0)的值(像素):[140, 3, 0, 0]
- 做一次 reverse:[0,0,3,140]
- 转成 int 类型:908
- 同样地,从纹理B中获取到整数 908
- 在 GPU 中计算 908*908
- 第二次循环计算 766*271
- 计算完成之后得到 sum。
- 5.转成像素
经过上面的计算,虽然可以得到 sum 值,但是 webgl 没法直接把值给到 js 进程,只能通过 canvas 获取到值,因此需要把结果转换成像素输出到 canvas。
以第一个计算结果为例:908908 + 766271 = 1032050。转换成 RGBA:
1111 10111111 01110010 -> [0,15,191,114]
此时获得的类型是:Uint8Array,实际是以 Uint32Array 表示一个整数,所以还需要做一次 reverse。
- 6.最终结果
Tensorflow.js 实现源码
传送门:tensorflow/tfjs
- 以乘法运算为例
const MUL = 'return a * b;';
export function multiply(){
...
if (env().getBool('WEBGL_PACK_BINARY_OPERATIONS')) {
program = new BinaryOpPackedProgram(MUL, a.shape, b.shape);
} else {
program = new BinaryOpProgram(MUL, a.shape, b.shape);
}
return backend.runWebGLProgram(program, [a, b], dtype);
}
- 先创建一个 program
export class BinaryOpPackedProgram implements GPGPUProgram {
constructor(op: string, aShape: number[], bShape: number[]){
this.userCode = `
vec4 binaryOperation(vec4 a, vec4 b) {
${op}
}
void main() {
vec4 a = getAAtOutCoords();
vec4 b = getBAtOutCoords();
vec4 result = binaryOperation(a, b);
${checkOutOfBoundsString}
setOutput(result);
}
`;
}
}
- 编译 program 程序
export class MathBackendWebGL extends KernelBackend {
gpgpu: GPGPUContext;
runWebGLProgram(){
...
const binary = this.getAndSaveBinary(key, () => {
return gpgpu_math.compileProgram(
this.gpgpu, program, inputsData, outputData);
});
...
}
}
- 创建片段着色器
export function compileProgram(){
...
const source = shader_compiler.makeShader(inputInfos, outShapeInfo, program);
const fragmentShader = createFragmentShader(gpgpu.gl, source);
const webGLProgram = gpgpu.createProgram(fragmentShader);
...
}
- 根据创建的片段着色器创建程序
export class GPGPUContext {
...
public createProgram(fragmentShader: WebGLShader): WebGLProgram {
this.throwIfDisposed();
const gl = this.gl;
// 片段着色器统一
if (this.vertexShader == null) {
this.vertexShader = gpgpu_util.createVertexShader(gl);
}
const program: WebGLProgram = webgl_util.createProgram(gl);
webgl_util.callAndCheck(
gl, () => gl.attachShader(program, this.vertexShader));
webgl_util.callAndCheck(gl, () => gl.attachShader(program, fragmentShader));
webgl_util.linkProgram(gl, program);
if (this.debug) {
webgl_util.validateProgram(gl, program);
}
if (!this.vertexAttrsAreBound) {
this.setProgram(program);
this.vertexAttrsAreBound = gpgpu_util.bindVertexProgramAttributeStreams(
gl, this.program, this.vertexBuffer);
}
return program;
}
...
}
WebGPU
WebGPU 被成为下一代 web 图形 API:
- WebGPU 是新的图形 API,将 GPU 硬件能力暴露给 Web 端。
- W3C WebGPU 社区团队在 2017Q1 成立,几乎所有的浏览器厂商都在组织内。
我们知道显卡都需要安装显卡驱动程序,通过显卡驱动程序暴露的 API 就可以操作 GPU。但是显卡驱动程序很难写,于是就出现了 OpenGL,专门负责上层接口封装并与下层显卡驱动打交道。
WebGL 2.0 是目前 WebGL 最高的版本,但这已经是 2017 年的标准,对应的 GPU Native 移动端的 OpenGL ES3.0 协议,这是 2012 年的标准,对应桌面端的协议是 OpenGL 3.2,这是 2009 年的标准。
WebGL 的 API 适应早期的 GPU 架构,没有函数编写能力,主要是一组状态机,用户通过传递参数来配置,实现一些规定的逻辑功能。而现在的 GPU 能力很强大,具备可编程能力,用户可以实现自定义的算法。
Microsoft 为此做出来最新的图形API 是 Direct3D 12,Apple 为此做出来最新的图形API 是 Metal,有一个著名的组织则做出来 Vulkan,这个组织名叫 Khronos。D3D12 现在在发光发热的地方是 Windows 和 PlayStation,Metal 则是 Mac 和 iPhone,Vulkan 你可能在安卓手机评测上见得多。这三个图形 API 被称作三大现代图形 API,与现代显卡(无论是PC还是移动设备)的联系很密切。
既然现在各个操作系统基本都没怎么装 Open GL了,那基于 OpenGL ES 的 WebGL 为什么能跑在各个操作系统的浏览器?因为目前基本都是靠三大现代图形 API 转译的。
除了在架构和性能方面,代码编写方面 WebGPU 也要比 WebGL 友好很多。在上面的代码中会经常出现 gl.xxx,每一次调用 gl.xxx,都需要完成 CPU 到 GPU 的信号传递,改变 GPU 的状态是立即生效的。毕竟原来在设计 GPU 的时候并没有考虑要干那么多事。而现代图形 API 更倾向于先把所有的东西准备好,然后通过一次性提交给 GPU。
WebGPU 目前还处于测试阶段,但是 tensorflow.js、three.js、babylon.js 都在实现对 WebGPU 的支持。
更多关于 WebGPU 和 WebGL 的差别可以参考:WebGPU,下一代 Web 图形技术。
其他
本文 demo 代码:webgl-demo。