add gpu call method for TimeDistributed wrapper layer

Author: transcranial

notebooks/layers/wrappers/TimeDistributed.ipynb

@@ -71,7 +71,7 @@
     "postcss-loader": "^2.0.8",
     "raw-loader": "^0.5.1",
     "vue": "^2.5.3",
-    "vue-loader": "^13.4.0",
+    "vue-loader": "^13.5.0",
     "vue-mdl": "^1.1.1",
     "vue-router": "^3.0.1",
     "vue-template-compiler": "^2.5.3",

package.json

@@ -36,6 +36,10 @@ export default class Bidirectional extends Layer {
     this.forwardLayer = new recurrentLayers[layer.class_name](forwardLayerAttrs)
     this.backwardLayer = new recurrentLayers[layer.class_name](backwardLayerAttrs)
 
+    // prevent GPU -> CPU data transfer by specifying non-empty outbound nodes array on internal layers
+    this.forwardLayer.outbound = [null]
+    this.backwardLayer.outbound = [null]
+
     this.mergeMode = merge_mode
     this.returnSequences = layer.config.return_sequences
 
@@ -148,6 +152,7 @@ export default class Bidirectional extends Layer {
       inputs: [{ texture: x.glTexture, type: '2d', name: 'source' }]
     })
 
+    // run internal component layers
     this.forwardLayer._callGPU(x)
     this.backwardLayer._callGPU(this.inputCopy)
     const forwardOutput = this.forwardLayer.output

src/layers/wrappers/Bidirectional.js

@@ -0,0 +1,21 @@
+#version 300 es
+precision highp float;
+
+in vec2 outTex;
+uniform sampler2D outputCopy;
+uniform sampler2D sliceOutput;
+uniform int t;
+uniform int timesteps;
+out vec4 outColor;
+
+void main() {
+  ivec2 size = textureSize(sliceOutput, 0);
+  int out_x = int(float(size[0]) * outTex.x);
+  int out_y = int(float(timesteps) * outTex.y);
+
+  if (t == out_y) {
+    outColor = vec4(texelFetch(sliceOutput, ivec2(out_x, 0), 0).r);
+  } else {
+    outColor = texelFetch(outputCopy, ivec2(out_x, out_y), 0);
+  }
+}

src/layers/wrappers/TimeDistributed.copySliceOutput.glsl

@@ -1,6 +1,8 @@
 import Layer from '../../Layer'
 import Tensor from '../../Tensor'
+import { webgl2 } from '../../WebGL2'
 import ops from 'ndarray-ops'
+import * as layers from '../'
 
 /**
  * TimeDistributed wrapper layer class
@@ -17,8 +19,24 @@ export default class TimeDistributed extends Layer {
 
     const { layer } = attrs
 
-    if (!layer) this.throwError('wrapped layer is undefined.')
-    this.layer = layer
+    if (!layer) {
+      this.throwError('wrapped layer is undefined.')
+    }
+
+    const wrappedLayerAttrs = Object.assign({}, layer.config, { gpu: attrs.gpu })
+    this.wrappedLayer = new layers[layer.class_name](wrappedLayerAttrs)
+
+    // prevent GPU -> CPU data transfer by specifying non-empty outbound nodes array on internal layer
+    this.wrappedLayer.outbound = [null]
+
+    // GPU setup
+    if (this.gpu) {
+      this.copyTextureProgram = webgl2.compileProgram(require('../../copyTexture.glsl'))
+      this.mapInputProgram = webgl2.compileProgram(require('../../mapInput.glsl'))
+      this.selectSliceProgram = webgl2.compileProgram(require('./TimeDistributed.selectSlice.glsl'))
+      this.copySliceOutputProgram = webgl2.compileProgram(require('./TimeDistributed.copySliceOutput.glsl'))
+      this.mapSliceOutputProgram = webgl2.compileProgram(require('./TimeDistributed.mapSliceOutput.glsl'))
+    }
   }
 
   /**
@@ -28,32 +46,282 @@ export default class TimeDistributed extends Layer {
    * @param {Tensor[]} weightsArr - array of weights which are instances of Tensor
    */
   setWeights(weightsArr) {
-    this.layer.setWeights(weightsArr)
+    this.wrappedLayer.setWeights(weightsArr)
   }
 
   /**
-   * Method for layer computational logic
+   * Layer computational logic
    *
    * @param {Tensor} x
    * @returns {Tensor}
    */
   call(x) {
-    const xStepShape = [...x.tensor.shape.slice(1)]
-    let xStep = new Tensor([], xStepShape)
-    ops.assign(xStep.tensor, x.tensor.pick(0, ...xStepShape.map(s => null)))
-    let yStep = this.layer.call(xStep)
-    const yStepShape = yStep.tensor.shape.slice()
-    let y = new Tensor([], [x.tensor.shape[0], ...yStepShape])
-    ops.assign(y.tensor.pick(0, ...yStepShape.map(s => null)), yStep.tensor)
-
-    for (let i = 1, steps = x.tensor.shape[0]; i < steps; i++) {
-      let xStep = new Tensor([], xStepShape)
-      ops.assign(xStep.tensor, x.tensor.pick(i, ...xStepShape.map(s => null)))
-      yStep = this.layer.call(xStep)
-      ops.assign(y.tensor.pick(i, ...yStepShape.map(s => null)), yStep.tensor)
-    }
-
-    x.tensor = y.tensor
-    return x
+    if (this.gpu) {
+      this._callGPU(x)
+    } else {
+      this._callCPU(x)
+    }
+    return this.output
+  }
+
+  /**
+   * CPU call
+   *
+   * @param {Tensor} x
+   */
+  _callCPU(x) {
+    const stepShape = [...x.tensor.shape.slice(1)]
+    const step = new Tensor([], stepShape)
+    ops.assign(step.tensor, x.tensor.pick(0, ...Array(stepShape.length).fill(null)))
+    let stepOutput = this.wrappedLayer.call(step)
+    const stepOutputShape = stepOutput.tensor.shape.slice()
+    this.output = new Tensor([], [x.tensor.shape[0], ...stepOutputShape])
+    ops.assign(this.output.tensor.pick(0, ...Array(stepOutputShape.length).fill(null)), stepOutput.tensor)
+    for (let i = 1, timesteps = x.tensor.shape[0]; i < timesteps; i++) {
+      ops.assign(step.tensor, x.tensor.pick(i, ...Array(stepShape.length).fill(null)))
+      stepOutput = this.wrappedLayer.call(step)
+      ops.assign(this.output.tensor.pick(i, ...Array(stepOutputShape.length).fill(null)), stepOutput.tensor)
+    }
+  }
+
+  /**
+   * Creates row/col index mappings to map input texture to time-distributed slices
+   *
+   * @param {Object} indicesForReshaped
+   */
+  _createIndexMap(indicesForReshaped) {
+    if (this.rowIndexMaps && this.colIndexMaps) {
+      return
+    }
+
+    const indicesRow = new Tensor(indicesForReshaped.row.data, indicesForReshaped.row.shape, { type: Int32Array })
+    const indicesCol = new Tensor(indicesForReshaped.col.data, indicesForReshaped.col.shape, { type: Int32Array })
+
+    this.rowIndexMaps = []
+    this.colIndexMaps = []
+
+    const timesteps = this.inputShape[0]
+    const sliceShape = this.inputShape.slice(1)
+    for (let t = 0; t < timesteps; t++) {
+      const sliceIndicesRow = new Tensor([], sliceShape, { type: Int32Array })
+      const sliceIndicesCol = new Tensor([], sliceShape, { type: Int32Array })
+      ops.assign(sliceIndicesRow.tensor, indicesRow.tensor.pick(t, ...Array(sliceShape.length).fill(null)))
+      ops.assign(sliceIndicesCol.tensor, indicesCol.tensor.pick(t, ...Array(sliceShape.length).fill(null)))
+      sliceIndicesRow.reshapeTo2DSquare()
+      sliceIndicesCol.reshapeTo2DSquare()
+      sliceIndicesRow.createGLTexture('2d', 'int')
+      sliceIndicesCol.createGLTexture('2d', 'int')
+      this.rowIndexMaps.push(sliceIndicesRow)
+      this.colIndexMaps.push(sliceIndicesCol)
+    }
+  }
+
+  /**
+   * Creates row/col index mappings to map time-distributed slices to output texture
+   *
+   * @param {Object} indicesForReshaped
+   */
+  _createOutputIndexMap(indicesForReshaped) {
+    if (this.outputRowIndexMaps && this.outputColIndexMaps) {
+      return
+    }
+
+    const outputSliceIndicesRow = new Tensor(indicesForReshaped.row.data, indicesForReshaped.row.shape, {
+      type: Int32Array
+    })
+    const outputSliceIndicesCol = new Tensor(indicesForReshaped.col.data, indicesForReshaped.col.shape, {
+      type: Int32Array
+    })
+
+    this.outputRowIndexMaps = []
+    this.outputColIndexMaps = []
+
+    const timesteps = this.outputShape[0]
+    const sliceShape = this.outputShape.slice(1)
+    for (let t = 0; t < timesteps; t++) {
+      const outputIndicesRow = new Tensor([], this.outputShape, { type: Int32Array })
+      const outputIndicesCol = new Tensor([], this.outputShape, { type: Int32Array })
+      ops.assigns(outputIndicesRow.tensor, -1)
+      ops.assigns(outputIndicesCol.tensor, -1)
+      ops.assign(outputIndicesRow.tensor.pick(t, ...Array(sliceShape.length).fill(null)), outputSliceIndicesRow.tensor)
+      ops.assign(outputIndicesCol.tensor.pick(t, ...Array(sliceShape.length).fill(null)), outputSliceIndicesCol.tensor)
+      outputIndicesRow.reshapeTo2DSquare()
+      outputIndicesCol.reshapeTo2DSquare()
+      outputIndicesRow.createGLTexture('2d', 'int')
+      outputIndicesCol.createGLTexture('2d', 'int')
+      this.outputRowIndexMaps.push(outputIndicesRow)
+      this.outputColIndexMaps.push(outputIndicesCol)
+    }
+  }
+
+  /**
+   * GPU call
+   *
+   * @param {Tensor} x
+   */
+  _callGPU(x) {
+    if (x.is2DReshaped) {
+      this.inputShape = x.originalShape
+    } else {
+      this.inputShape = x.tensor.shape
+    }
+
+    if (!x.glTexture) {
+      if (x.tensor.shape.length <= 2) {
+        x.createGLTexture()
+      } else if (x.tensor.shape.length > 2 && !x.is2DReshaped) {
+        x.reshapeTo2DSquare()
+        x.createGLTexture()
+      }
+    }
+
+    if (this.inputShape.length > 2) {
+      this._createIndexMap(x.indicesForReshaped)
+    }
+
+    const timesteps = this.inputShape[0]
+    const sliceShape = this.inputShape.slice(1)
+
+    if (!this.slice) {
+      this.slice = new Tensor([], sliceShape)
+      if (sliceShape.length <= 2) {
+        this.slice.createGLTexture()
+      } else {
+        this.slice.reshapeTo2DSquare()
+        this.slice.createGLTexture()
+      }
+    }
+
+    if (this.inputShape.length <= 2) {
+      webgl2.runProgram({
+        program: this.selectSliceProgram,
+        output: this.slice,
+        inputs: [{ texture: x.glTexture, type: '2d', name: 'x' }],
+        uniforms: [{ value: 0, type: 'int', name: 't' }]
+      })
+    } else {
+      webgl2.runProgram({
+        program: this.mapInputProgram,
+        output: this.slice,
+        inputs: [
+          { texture: x.glTexture, type: '2d', name: 'x' },
+          { texture: this.rowIndexMaps[0].glTexture, type: '2d', name: 'rowIndexMap' },
+          { texture: this.colIndexMaps[0].glTexture, type: '2d', name: 'colIndexMap' }
+        ]
+      })
+    }
+
+    this.wrappedLayer._callGPU(this.slice)
+    this.sliceOutput = this.wrappedLayer.output
+
+    if (!this.output) {
+      if (this.inputShape.length <= 2) {
+        this.outputShape = [timesteps, this.sliceOutput.glTextureShape[1]]
+        this.output = new Tensor([], this.outputShape)
+        this.outputCopy = new Tensor([], this.outputShape)
+        this.output.createGLTexture()
+        this.outputCopy.createGLTexture()
+      } else {
+        this.outputShape = [timesteps, ...this.sliceOutput.originalShape]
+        this.output = new Tensor([], this.outputShape)
+        this.outputCopy = new Tensor([], this.outputShape)
+        this.output.reshapeTo2DSquare()
+        this.outputCopy.reshapeTo2DSquare()
+        this.output.createGLTexture()
+        this.outputCopy.createGLTexture()
+
+        this._createOutputIndexMap(this.sliceOutput.indicesForReshaped)
+      }
+    }
+
+    webgl2.runProgram({
+      program: this.copyTextureProgram,
+      output: this.outputCopy,
+      inputs: [{ texture: this.output.glTexture, type: '2d', name: 'source' }]
+    })
+
+    if (this.inputShape.length <= 2) {
+      webgl2.runProgram({
+        program: this.copySliceOutputProgram,
+        output: this.output,
+        inputs: [
+          { texture: this.outputCopy.glTexture, type: '2d', name: 'outputCopy' },
+          { texture: this.sliceOutput.glTexture, type: '2d', name: 'sliceOutput' }
+        ],
+        uniforms: [{ value: 0, type: 'int', name: 't' }, { value: timesteps, type: 'int', name: 'timesteps' }]
+      })
+    } else {
+      webgl2.runProgram({
+        program: this.mapSliceOutputProgram,
+        output: this.output,
+        inputs: [
+          { texture: this.outputCopy.glTexture, type: '2d', name: 'outputCopy' },
+          { texture: this.sliceOutput.glTexture, type: '2d', name: 'sliceOutput' },
+          { texture: this.outputRowIndexMaps[0].glTexture, type: '2d', name: 'rowIndexMap' },
+          { texture: this.outputColIndexMaps[0].glTexture, type: '2d', name: 'colIndexMap' }
+        ]
+      })
+    }
+
+    for (let i = 1; i < timesteps; i++) {
+      if (this.inputShape.length <= 2) {
+        webgl2.runProgram({
+          program: this.selectSliceProgram,
+          output: this.slice,
+          inputs: [{ texture: x.glTexture, type: '2d', name: 'x' }],
+          uniforms: [{ value: i, type: 'int', name: 't' }]
+        })
+      } else {
+        webgl2.runProgram({
+          program: this.mapInputProgram,
+          output: this.slice,
+          inputs: [
+            { texture: x.glTexture, type: '2d', name: 'x' },
+            { texture: this.rowIndexMaps[i].glTexture, type: '2d', name: 'rowIndexMap' },
+            { texture: this.colIndexMaps[i].glTexture, type: '2d', name: 'colIndexMap' }
+          ]
+        })
+      }
+
+      this.wrappedLayer._callGPU(this.slice)
+      this.sliceOutput = this.wrappedLayer.output
+
+      webgl2.runProgram({
+        program: this.copyTextureProgram,
+        output: this.outputCopy,
+        inputs: [{ texture: this.output.glTexture, type: '2d', name: 'source' }]
+      })
+
+      if (this.inputShape.length <= 2) {
+        webgl2.runProgram({
+          program: this.copySliceOutputProgram,
+          output: this.output,
+          inputs: [
+            { texture: this.outputCopy.glTexture, type: '2d', name: 'outputCopy' },
+            { texture: this.sliceOutput.glTexture, type: '2d', name: 'sliceOutput' }
+          ],
+          uniforms: [{ value: i, type: 'int', name: 't' }, { value: timesteps, type: 'int', name: 'timesteps' }]
+        })
+      } else {
+        webgl2.runProgram({
+          program: this.mapSliceOutputProgram,
+          output: this.output,
+          inputs: [
+            { texture: this.outputCopy.glTexture, type: '2d', name: 'outputCopy' },
+            { texture: this.sliceOutput.glTexture, type: '2d', name: 'sliceOutput' },
+            { texture: this.outputRowIndexMaps[i].glTexture, type: '2d', name: 'rowIndexMap' },
+            { texture: this.outputColIndexMaps[i].glTexture, type: '2d', name: 'colIndexMap' }
+          ]
+        })
+      }
+    }
+
+    // GPU -> CPU data transfer
+    if (this.outbound.length === 0) {
+      this.output.transferFromGLTexture()
+      if (this.output.is2DReshaped) {
+        this.output.reshapeFrom2DSquare()
+      }
+    }
   }
 }