graceful fallback to CPU when tensor MAX_TEXTURE_SIZE exceeded (#8)

src/Tensor.js

@@ -31,6 +31,8 @@ export default class Tensor {
     } else {
       this.tensor = ndarray(new this._type([]), [])
     }
+
+    this._gpuMaxSizeExceeded = false
   }
 
   /**
@@ -39,60 +41,24 @@ export default class Tensor {
    * see https://github.com/waylonflinn/weblas/wiki/Pipeline
    *
    * gl.MAX_TEXTURE_SIZE is a limiting factor.
-   * Where this is exceeded, weblas Tensor must be split.
+   * Where this is exceeded, falls back to CPU.
    */
   createWeblasTensor () {
     if (this.weblasTensor) {
       this.weblasTensor.delete()
     }
-    if (this.weblasTensorsSplit) {
-      this.weblasTensorsSplit.forEach(t => t.delete())
-    }
 
     if (this.tensor.shape.length === 1) {
       const len = this.tensor.shape[0]
       if (len > MAX_TEXTURE_SIZE) {
-        this.weblasTensorsSplit = []
-        const splitNum = Math.ceil(MAX_TEXTURE_SIZE / len)
-        for (let i = 0; i < splitNum; i++) {
-          const lo = i * Math.round(len / splitNum)
-          const hi = Math.min(len, (i + 1) * Math.round(len / splitNum))
-          const splitShape = [1, hi - lo]
-          this.weblasTensorsSplit.push(new weblas.pipeline.Tensor(splitShape, this.tensor.data.subarray(lo, hi - lo)))
-        }
+        this._gpuMaxSizeExceeded = true
       } else {
         const shape = [1, len]
         this.weblasTensor = new weblas.pipeline.Tensor(shape, this.tensor.data)
       }
     } else if (this.tensor.shape.length === 2) {
-      if (this.tensor.shape.every(s => s > MAX_TEXTURE_SIZE)) {
-        throw new Error('[Tensor] cannot create Tensor with both dimensions exceeding MAX_TEXTURE_SIZE')
-      }
-
-      const rows = this.tensor.shape[0]
-      const cols = this.tensor.shape[1]
-      if (rows > MAX_TEXTURE_SIZE) {
-        this.weblasTensorsSplit = []
-        const splitNum = Math.ceil(rows / MAX_TEXTURE_SIZE)
-        for (let i = 0; i < splitNum; i++) {
-          const lo = i * MAX_TEXTURE_SIZE
-          const hi = Math.min(rows, (i + 1) * MAX_TEXTURE_SIZE)
-          const splitShape = [hi - lo, cols]
-          this.weblasTensorsSplit.push(
-            new weblas.pipeline.Tensor(splitShape, this.tensor.data.subarray(lo * cols, hi * cols))
-          )
-        }
-      } else if (cols > MAX_TEXTURE_SIZE) {
-        this.weblasTensorsSplit = []
-        const splitNum = Math.ceil(cols / MAX_TEXTURE_SIZE)
-        for (let i = 0; i < splitNum; i++) {
-          const lo = i * MAX_TEXTURE_SIZE
-          const hi = Math.min(cols, (i + 1) * MAX_TEXTURE_SIZE)
-          const splitShape = [rows, hi - lo]
-          this.weblasTensorsSplit.push(
-            new weblas.pipeline.Tensor(splitShape, this.tensor.data.slice(rows * lo, rows * hi))
-          )
-        }
+      if (this.tensor.shape.some(s => s > MAX_TEXTURE_SIZE)) {
+        this._gpuMaxSizeExceeded = true
       } else {
         const shape = this.tensor.shape
         this.weblasTensor = new weblas.pipeline.Tensor(shape, this.tensor.data)

src/layers/convolutional/Convolution2D.js

@@ -167,14 +167,12 @@ export default class Convolution2D extends Layer {
     }
 
     let patch = new Tensor([], [nbRow, nbCol, inputChannels])
-    let patchRaveled = new Tensor([], [patchLen])
-    let n = 0
+    let offset = 0
     for (let i = 0, limit = inputRows - nbRow; i <= limit; i += this.subsample[0]) {
       for (let j = 0, limit = inputCols - nbCol; j <= limit; j += this.subsample[1]) {
         ops.assign(patch.tensor, x.tensor.hi(i + nbRow, j + nbCol, inputChannels).lo(i, j, 0))
-        patchRaveled.replaceTensorData(patch.tensor.data)
-        ops.assign(this._imColsMat.tensor.pick(n, null), patchRaveled.tensor)
-        n += 1
+        this._imColsMat.tensor.data.set(patch.tensor.data, offset)
+        offset += patchLen
       }
     }
     if (this._useWeblas) {
@@ -227,27 +225,13 @@ export default class Convolution2D extends Layer {
     const nbPatches = outputRows * outputCols
     const matMul = new Tensor([], [nbPatches, nbFilter])
 
-    if (this._useWeblas) {
+    if (this._useWeblas && !(this._imColsMat._gpuMaxSizeExceeded || this._wRowsMat._gpuMaxSizeExceeded)) {
       // GPU
       const bias = this.bias ? this.weights.b.weblasTensor : this._zerosVec.weblasTensor
-      if (this._imColsMat.weblasTensorsSplit) {
-        // split matrix multiply if this._imColsMat dimension > webgl.MAX_TEXTURE_SIZE
-        let offset = 0
-        this._imColsMat.weblasTensorsSplit.forEach(imColsMatSplit => {
-          const matMulSplitData = weblas.pipeline.sgemm(
-            1, imColsMatSplit, this._wRowsMat.weblasTensor,
-            1, bias
-          ).transfer()
-          matMul.tensor.data.set(matMulSplitData, offset)
-          offset += matMulSplitData.length
-        })
-      } else {
-        // normal matrix multiply
-        matMul.tensor.data = weblas.pipeline.sgemm(
-          1, this._imColsMat.weblasTensor, this._wRowsMat.weblasTensor,
-          1, bias
-        ).transfer()
-      }
+      matMul.tensor.data = weblas.pipeline.sgemm(
+        1, this._imColsMat.weblasTensor, this._wRowsMat.weblasTensor,
+        1, bias
+      ).transfer()
     } else {
       // CPU
       if (this.bias) {