使用 tensorflow.js v3.13.0 进行模型推理后，如何使用保存在 GPU 上的数据？

Question

我刚刚了解到 tfjs v3.13.0 中的一项新功能（请参阅 https://github.com/tensorflow/tfjs/pull/5953）。我正在尝试使用 Tensors 的新 dataToGPU() 方法将我的模型输出保存在 GPU 上，因为将数据发送回 CPU 的 data() 方法对我来说花费了太多时间用例。但是当我调用新方法并尝试将它创建的 WebGLTexture 绑定到我的 WebGLRenderingContext 时，我收到以下错误。

WebGL: INVALID_OPERATION: bindTexture: object does not belong to this context

我猜这是因为纹理是在与我要绑定纹理的 canvas 不同的上下文中创建的。因此，为了解决这个问题，似乎还有另一个功能，它增加了向 tfjs 的 WebGL 后端声明提供 HTMLCanvasElement 或 OffscreenCanvas 的可能性（参见 https://github.com/tensorflow/tfjs/pull/5983）。但是，我没有在我的代码中声明任何后端，所以我不确定我必须如何使用这些功能。

有人可以告诉我在运行模型时应该如何实例化和使用 WebGL 后端吗？

Answer 1

有关如何注册基于 webgl 的自定义后端的示例，请参阅 GitHub https://github.com/vladmandic/human/blob/main/src/tfjs/humangl.ts

上的以下内容

此处添加此代码的副本以防上述 link 失败：

/** TFJS custom backend registration */

import type { Human } from '../human';
import { log } from '../util/util';
import * as tf from '../../dist/tfjs.esm.js';
import * as image from '../image/image';
import * as models from '../models';
import type { AnyCanvas } from '../exports';
// import { env } from '../env';

export const config = {
  name: 'humangl',
  priority: 999,
  canvas: <null | AnyCanvas>null,
  gl: <null | WebGL2RenderingContext>null,
  extensions: <string[]> [],
  webGLattr: { // https://www.khronos.org/registry/webgl/specs/latest/1.0/#5.2
    alpha: false,
    antialias: false,
    premultipliedAlpha: false,
    preserveDrawingBuffer: false,
    depth: false,
    stencil: false,
    failIfMajorPerformanceCaveat: false,
    desynchronized: true,
  },
};

function extensions(): void {
  /*
  https://www.khronos.org/registry/webgl/extensions/
  https://webglreport.com/?v=2
  */
  const gl = config.gl;
  if (!gl) return;
  config.extensions = gl.getSupportedExtensions() as string[];
  // gl.getExtension('KHR_parallel_shader_compile');
}

/**
 * Registers custom WebGL2 backend to be used by Human library
 *
 * @returns void
 */
export async function register(instance: Human): Promise<void> {
  // force backend reload if gl context is not valid
  if (instance.config.backend !== 'humangl') return;
  if ((config.name in tf.engine().registry) && (!config.gl || !config.gl.getParameter(config.gl.VERSION))) {
    log('error: humangl backend invalid context');
    models.reset(instance);
    /*
    log('resetting humangl backend');
    await tf.removeBackend(config.name);
    await register(instance); // re-register
    */
  }
  if (!tf.findBackend(config.name)) {
    try {
      config.canvas = await image.canvas(100, 100);
    } catch (err) {
      log('error: cannot create canvas:', err);
      return;
    }
    try {
      config.gl = config.canvas?.getContext('webgl2', config.webGLattr) as WebGL2RenderingContext;
      const glv2 = config.gl.getParameter(config.gl.VERSION).includes('2.0');
      if (!glv2) {
        log('override: using fallback webgl backend as webgl 2.0 is not detected');
        instance.config.backend = 'webgl';
        return;
      }
      if (config.canvas) {
        config.canvas.addEventListener('webglcontextlost', async (e) => {
          log('error: humangl:', e.type);
          log('possible browser memory leak using webgl or conflict with multiple backend registrations');
          instance.emit('error');
          throw new Error('backend error: webgl context lost');
          // log('resetting humangl backend');
          // env.initial = true;
          // models.reset(instance);
          // await tf.removeBackend(config.name);
          // await register(instance); // re-register
        });
        config.canvas.addEventListener('webglcontextrestored', (e) => {
          log('error: humangl context restored:', e);
        });
        config.canvas.addEventListener('webglcontextcreationerror', (e) => {
          log('error: humangl context create:', e);
        });
      }
    } catch (err) {
      log('error: cannot get WebGL context:', err);
      return;
    }
    try {
      tf.setWebGLContext(2, config.gl);
    } catch (err) {
      log('error: cannot set WebGL context:', err);
      return;
    }
    try {
      const ctx = new tf.GPGPUContext(config.gl);
      tf.registerBackend(config.name, () => new tf.MathBackendWebGL(ctx), config.priority);
    } catch (err) {
      log('error: cannot register WebGL backend:', err);
      return;
    }
    try {
      const kernels = tf.getKernelsForBackend('webgl');
      kernels.forEach((kernelConfig) => {
        const newKernelConfig = { ...kernelConfig, backendName: config.name };
        tf.registerKernel(newKernelConfig);
      });
    } catch (err) {
      log('error: cannot update WebGL backend registration:', err);
      return;
    }
    const current = tf.backend().getGPGPUContext ? tf.backend().getGPGPUContext().gl : null;
    if (current) {
      log(`humangl webgl version:${current.getParameter(current.VERSION)} renderer:${current.getParameter(current.RENDERER)}`);
    } else {
      log('error: no current gl context:', current, config.gl);
      return;
    }
    try {
      tf.ENV.set('WEBGL_VERSION', 2);
    } catch (err) {
      log('error: cannot set WebGL backend flags:', err);
      return;
    }
    extensions();
    log('backend registered:', config.name);
  }
}

使用 tensorflow.js v3.13.0 进行模型推理后，如何使用保存在 GPU 上的数据？

How to use data kept on GPU after model inference with tensorflow.js v3.13.0?

gpu

shader

webgl

tensorflow.js