TensorFlowJS之DataStorage的set/get和write/read

作者的经验是，set/get其实就是对对象数据的读和写。TypeScript实现的DataStorage看起来似乎有点另类（作者第一次分析这么庞大的JavaScript/TypeScript项目）。另类在：它既有set/get，同时还有write/read。

Set/Get

set/get是直接封装在DataStorage的定义里面的：

  write(dataId: DataId, values: DataValues): void {
    if (values == null) {
      throw new Error('MathBackendCPU.write(): values can not be null');
    }
    this.data.get(dataId).values = values;
  }
  async read(dataId: DataId): Promise<DataValues> {
    return this.readSync(dataId);
  }
  readSync(dataId: DataId): DataValues {
    const {dtype, complexTensors} = this.data.get(dataId);
    if (dtype === 'complex64') {
      const realValues = complexTensors.real.dataSync() as Float32Array;
      const imagValues = complexTensors.imag.dataSync() as Float32Array;
      return complex_util.mergeRealAndImagArrays(realValues, imagValues);
    }
    return this.data.get(dataId).values;
  }  

Write/Read

write/read就没那么直接。在不同的backend里面实现了不同的write/read。

backend_cpu.ts实现的write/read：

  write(dataId: DataId, values: DataValues): void {
    if (values == null) {
      throw new Error('MathBackendCPU.write(): values can not be null');
    }
    this.data.get(dataId).values = values;
  }
  async read(dataId: DataId): Promise<DataValues> {
    return this.readSync(dataId);
  }
  readSync(dataId: DataId): DataValues {
    const {dtype, complexTensors} = this.data.get(dataId);
    if (dtype === 'complex64') {
      const realValues = complexTensors.real.dataSync() as Float32Array;
      const imagValues = complexTensors.imag.dataSync() as Float32Array;
      return complex_util.mergeRealAndImagArrays(realValues, imagValues);
    }
    return this.data.get(dataId).values;
  }  

backend_webgl.ts实现的write/read和cpu write/read类似的地方在于，数据写入的时候是写入到texData.values。读取的时候，则是从texData.values里面读取。区别在于：写入的时候，写缓存到texData.values，然后传递给GPU:

  private uploadToGPU(dataId: DataId): void {
    const texData = this.texData.get(dataId);
    const {shape, values, texture, usage, isPacked} = texData;

读取的时候，则是通过readPixels从GPU读取过来的：

  private getValuesFromTexture(dataId: DataId): Float32Array {
    const {shape, dtype, texture, texShape} = this.texData.get(dataId);
    const size = util.sizeFromShape(shape);
    if (ENV.get('WEBGL_DOWNLOAD_FLOAT_ENABLED')) {
      if (this.texData.get(dataId).isPacked) {
        const batch = webgl_util.getBatchDim(shape);
        let rows = 1, cols = 1;
        if (shape.length) {
          [rows, cols] = webgl_util.getRowsCols(shape);
        }
        return this.gpgpu
            .downloadMatrixFromPackedTexture(
                texture, batch, rows, cols, texShape[0], texShape[1])
            .subarray(0, size);
      } else {
        return this.gpgpu
            .downloadFloat32MatrixFromOutputTexture(
                texture, texShape[0], texShape[1])
            .subarray(0, size);
      }
    }

小结

set/get仅仅是用来管理多个对象的，是对象管理器。write/read则是为了操作对象里面的数据的，是对象本身的操作。