TensorFlowJS之Shader的生成

最新推荐文章于 2025-10-12 15:29:27 发布

原创最新推荐文章于 2025-10-12 15:29:27 发布 · 249 阅读

0 ·

CC 4.0 BY-SA版权

TensorFlowJS 专栏收录该内容

14 篇文章

订阅专栏

本文深入探讨了Shader如何根据用户输入的Tensor和输出的Tensor自动生成的过程。详细解析了makeShader函数如何遍历inputsInfo，为每个InputInfo生成相应的采样器和位置信息，并介绍了UV坐标在顶点着色器中的作用及输出变量的Shader生成方式。

Shader是根据用户输入的Tensor和输出的Tensor自动生成的。生成的函数是makeShader。

输入变量的Shader

输入包括采样器和UV坐标。
如果定义了两个Tensor，这段代码运行的结果是生成两个采样器：

uniform sampler2D A;
uniform int offsetA;
uniform sampler2D B;
uniform int offsetB;

makeShader会遍历inputsInfo，对每个InputInfo生成相应的采样器和位置信息（offset*是针对某些输入会被slice的情形）：

export function makeShader(
    inputsInfo: InputInfo[], outputShape: ShapeInfo, userCode: string,
    usesPackedTextures: boolean): string {
  const prefixSnippets: string[] = [];
  inputsInfo.forEach(x => {
    const size = util.sizeFromShape(x.shapeInfo.logicalShape);
    // Snippet when we decided to upload the values as uniform.
    if (x.shapeInfo.isUniform) {
      prefixSnippets.push(
          `uniform float ${x.name}${size > 1 ? `[${size}]` : ''};`);
    } else {
      prefixSnippets.push(`uniform sampler2D ${x.name};`);
      prefixSnippets.push(`uniform int offset${x.name};`);
    }
  });
  const inputPrefixSnippet = prefixSnippets.join('\n');

  const inputSamplingSnippet =
      inputsInfo
          .map(x => getInputSamplingSnippet(x, outputShape, usesPackedTextures))
          .join('\n');

UV坐标的生成（其实UV坐标是来自顶点着色器的，这里对顶点着色器传递过来的UV坐标进行了加工）：

function getOutputSamplingSnippet(
    outShape: number[], outTexShape: [number, number]): string {
  switch (outShape.length) {
    case 0:
      return getOutputScalarCoords();
    case 1:
      return getOutput1DCoords(outShape as [number], outTexShape);
    case 2:
      return getOutput2DCoords(outShape as [number, number], outTexShape);
    case 3:
      return getOutput3DCoords(
          outShape as [number, number, number], outTexShape);
    case 4:
      return getOutput4DCoords(
          outShape as [number, number, number, number], outTexShape);
    case 5:
      return getOutput5DCoords(
          outShape as [number, number, number, number, number], outTexShape);
    case 6:
      return getOutput6DCoords(
          outShape as [number, number, number, number, number, number],
          outTexShape);
    default:
      throw new Error(
          `${outShape.length}-D output sampling is not yet supported`);
  }
}

输出变量的Shader

通常只定义一个输出。

function getFloatTextureSetRSnippet(glsl: GLSL): string {
  return `
    void setOutput(float val) {
      ${glsl.output} = vec4(val, 0, 0, 0);
    }
  `;
}

对于一种情况，输出是：

void setOutput(float val) {
   outputColor = vec4(val, 0, 0, 0);
}

JavaScript-结构体的map-join

TFJS里面很常见这个map-join模式。

function updateItem(x) {
	return x.name+":"+x.value;
}
var array = [];
array.push({ name: 1, value: "a" });
array.push({ name: 2, value: "b" });
const map1 = array.map(x => updateItem(x));

function logMapElements(value, key, map) {
  console.log(`m[${key}] = ${value}`);
}

map1.forEach(logMapElements);

输出结果：

> "m[0] = 1:a"
> "m[1] = 2:b"
> "1:a,2:b"