slight im2col op speedup

src/layers/convolutional/AtrousConvolution2D.js

@@ -88,8 +88,7 @@ export default class AtrousConvolution2D extends Convolution2D {
     }
 
     let patch = new Tensor([], [nbRow, nbCol, inputChannels])
-    let patchRaveled = new Tensor([], [patchLen])
-    let n = 0
+    let offset = 0
     for (let i = 0, limit = inputRows - nbRowDilated; i <= limit; i += this.subsample[0]) {
       for (let j = 0, limit = inputCols - nbColDilated; j <= limit; j += this.subsample[1]) {
         ops.assign(
@@ -99,9 +98,8 @@ export default class AtrousConvolution2D extends Convolution2D {
             .lo(i, j, 0)
             .step(this.atrousRate[0], this.atrousRate[1], 1)
         )
-        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) {

src/layers/convolutional/Convolution3D.js

@@ -187,15 +187,13 @@ export default class Convolution3D extends Layer {
     }
 
     let patch = new Tensor([], [kernelDim1, kernelDim2, kernelDim3, inputChannels])
-    let patchRaveled = new Tensor([], [patchLen])
-    let n = 0
+    let offset = 0
     for (let i = 0, limit = inputDim1 - kernelDim1; i <= limit; i += this.subsample[0]) {
       for (let j = 0, limit = inputDim2 - kernelDim2; j <= limit; j += this.subsample[1]) {
         for (let k = 0, limit = inputDim3 - kernelDim3; k <= limit; k += this.subsample[2]) {
           ops.assign(patch.tensor, x.tensor.hi(i + kernelDim1, j + kernelDim2, k + kernelDim3, inputChannels).lo(i, j, k, 0))
-          patchRaveled.replaceTensorData(patch.tensor.data)
-          ops.assign(this._volColsMat.tensor.pick(n, null), patchRaveled.tensor)
-          n += 1
+          this._volColsMat.tensor.data.set(patch.tensor.data, offset)
+          offset += patchLen
         }
       }
     }

src/layers/convolutional/SeparableConvolution2D.js

@@ -32,18 +32,17 @@ class _DepthwiseConvolution2D extends Convolution2D {
     }
 
     let patch = new Tensor([], [nbRow, nbCol, 1])
-    let patchRaveled = new Tensor([], [patchLen])
-    let n = 0
+    let offset = 0
     for (let c = 0; c < inputChannels; c++) {
       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, c + 1).lo(i, j, c))
-          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) {
       this._imColsMat.createWeblasTensor()
     }
@@ -82,22 +81,30 @@ class _DepthwiseConvolution2D extends Convolution2D {
    * @returns {Tensor} x
    */
   call (x) {
-    // convert to tf ordering
-    if (this.dimOrdering === 'th') {
-      x.tensor = x.tensor.transpose(1, 2, 0)
-    }
-
+    let startTime = performance.now()
     this._calcOutputShape(x)
+    let endTime = performance.now()
+    console.log(0, endTime - startTime)
+    startTime = performance.now()
     this._padInput(x)
+    endTime = performance.now()
+    console.log(1, endTime - startTime)
 
+    startTime = performance.now()
     this._im2col(x)
+    endTime = performance.now()
+    console.log(2, endTime - startTime)
 
+    startTime = performance.now()
     const nbFilter = this.kernelShape[0]
     const outputRows = this.outputShape[0]
     const outputCols = this.outputShape[1]
     const nbPatches = outputRows * outputCols
     const matMul = new Tensor([], [nbPatches * x.tensor.shape[2], nbFilter * x.tensor.shape[2]])
+    endTime = performance.now()
+    console.log(3, endTime - startTime)
 
+    startTime = performance.now()
     if (this._useWeblas) {
       // GPU
       if (this._imColsMat.weblasTensorsSplit) {
@@ -122,7 +129,10 @@ class _DepthwiseConvolution2D extends Convolution2D {
       // CPU
       gemm(matMul.tensor, this._imColsMat.tensor, this._wRowsMat.tensor, 1, 1)
     }
+    endTime = performance.now()
+    console.log(4, endTime - startTime)
 
+    startTime = performance.now()
     let output = new Tensor([], [outputRows, outputCols, x.tensor.shape[2] * nbFilter])
     const outputDataLength = outputRows * outputCols * x.tensor.shape[2] * nbFilter
     let dataFiltered = new Float32Array(outputDataLength)
@@ -134,6 +144,8 @@ class _DepthwiseConvolution2D extends Convolution2D {
       }
     }
     output.replaceTensorData(dataFiltered)
+    endTime = performance.now()
+    console.log(5, endTime - startTime)
 
     x.tensor = output.tensor
 
@@ -221,7 +233,10 @@ export default class SeparableConvolution2D extends Layer {
     const depthwiseOutput = this._depthwiseConv.call(x)
 
     // Perform depthwise ops
+    let startTime = performance.now()
     const pointwiseOutput = this._pointwiseConv.call(depthwiseOutput)
+    let endTime = performance.now()
+    console.log(6, endTime - startTime)
 
     x.tensor = pointwiseOutput.tensor