Merge pull request #45 from trulyspinach/fix-reshape

Fix ONNX frontend Reshape behavior

Author: petrohi

tools/src/tensil/tools/compiler/OnnxFrontend.scala

@@ -975,11 +975,39 @@ class OnnxFrontend(
       throw new CompilerException("Scalar sizes must match for reshape")
 
     val indexesAndOffsetsPairs =
-      for (i <- 0 until inputDims.sizeScalars) yield {
-        val (inputIndex, inputOffset)   = inputDims.vectorIndexOffsetAt(i)
-        val (outputIndex, outputOffset) = outputDims.vectorIndexOffsetAt(i)
-        ((outputIndex, inputIndex), (outputOffset, inputOffset))
-      }
+      if (inputDims.order == 2 && outputDims.order == 4) {
+        // when reshaping from 1D to 4D, a NCHW to NHWC permute(transpose)
+        // need to be performed to match ONNX's behavior.
+
+        val dimMap    = (3 to 0 by -1).map(shape.takeRight(_).fold(1)(_ * _))
+        val outShape  = Seq(shape(0), shape(2), shape(3), shape(1))
+        val dimMapOut = (3 to 0 by -1).map(outShape.takeRight(_).fold(1)(_ * _))
+
+        for (
+          n <- 0 until shape(0);
+          c <- 0 until shape(1);
+          h <- 0 until shape(2);
+          w <- 0 until shape(3)
+        ) yield {
+
+          val from = Array(n, c, h, w).zipWithIndex
+            .map { case (e, i) => dimMap(i) * e }
+            .fold(0)(_ + _)
+
+          val to = Array(n, h, w, c).zipWithIndex
+            .map { case (e, i) => dimMapOut(i) * e }
+            .fold(0)(_ + _)
+
+          val (inputIndex, inputOffset)   = inputDims.vectorIndexOffsetAt(from)
+          val (outputIndex, outputOffset) = outputDims.vectorIndexOffsetAt(to)
+          ((outputIndex, inputIndex), (outputOffset, inputOffset))
+        }
+      } else
+        for (i <- 0 until inputDims.sizeScalars) yield {
+          val (inputIndex, inputOffset)   = inputDims.vectorIndexOffsetAt(i)
+          val (outputIndex, outputOffset) = outputDims.vectorIndexOffsetAt(i)
+          ((outputIndex, inputIndex), (outputOffset, inputOffset))
+        }
 
     val groupedByOffsetPairs = indexesAndOffsetsPairs
       .groupBy(_._1)

tools/test/src/tools/CompilerSpec.scala

@@ -1251,6 +1251,23 @@ class CompilerSpec extends FlatSpec {
     localDepth = 256,
   )
 
+  it should "Compile ONNX Reshape from 1D(NCHW) to 4D(NHWC)" in {
+    val name = "reshape_1d_4d"
+    val options = CompilerOptions(
+      arch = Conv2DTiny4x4Architecure,
+      printSummary = true
+    )
+
+    Compiler.compile(
+      name,
+      s"${Models}/reshape_1d_4d.onnx",
+      List("output"),
+      options
+    )
+
+    GoldenProcessorHelper.test(name, inputBatchSize = options.inputBatchSize)
+  }
+
   it should "Compile TF Conv2D (VALID padding) 3x3x4 image with 2x2x4x4 kernel" in {
     val name = "conv2d_4x4_valid"
     val options = CompilerOptions(

tools/test/src/tools/GoldenProcessorHelper.scala

@@ -95,6 +95,8 @@ object GoldenProcessorHelper {
       ResNet.prepareInputStream(dataType, arraySize, count)
     else if (modelName.startsWith("resnet50v2"))
       ResNet50.prepareInputStream(dataType, arraySize, count)
+    else if (modelName.startsWith("reshape_1d_4d"))
+      Reshape.prepareInputStream(dataType, arraySize, count)
     else if (yoloPattern.findFirstIn(modelName).isDefined) {
       val yoloPattern(yoloSize) = modelName
       TinyYolo(yoloSize.toInt, onnx = modelName.endsWith("onnx"))
@@ -174,6 +176,8 @@ object GoldenProcessorHelper {
       ResNet.assertOutput(dataType, arraySize, bytes, count)
     else if (modelName.startsWith("resnet50v2"))
       ResNet50.assertOutput(dataType, arraySize, bytes, count)
+    else if (modelName.startsWith("reshape_1d_4d"))
+      Reshape.assertOutput(dataType, arraySize, bytes, count)
     else if (yoloPattern.findFirstIn(modelName).isDefined) {
       val yoloPattern(yoloSize) = modelName
       TinyYolo(yoloSize.toInt, onnx = modelName.endsWith("onnx"))

tools/test/src/tools/Reshape.scala

@@ -0,0 +1,67 @@
+/* SPDX-License-Identifier: Apache-2.0 */
+/* Copyright © 2019-2022 Tensil AI Company */
+
+package tensil.tools
+
+import java.io._
+import scala.reflect.ClassTag
+import tensil.tools.golden.{Processor, ExecutiveTraceContext}
+import scala.collection.mutable
+import tensil.ArchitectureDataType
+
+object Reshape {
+  def prepareInputStream(
+      dataType: ArchitectureDataType,
+      arraySize: Int,
+      count: Int = 1
+  ): InputStream =
+    new ByteArrayInputStream(prepareInputBytes(dataType, arraySize, count))
+
+  def prepareInputBytes(
+      dataType: ArchitectureDataType,
+      arraySize: Int,
+      count: Int = 1
+  ): Array[Byte] = {
+    val inputPrep           = new ByteArrayOutputStream()
+    val inputPrepDataStream = new DataOutputStream(inputPrep)
+
+    val seq = (1 to 8).map(_.toFloat).toArray.grouped(arraySize)
+    for (s <- seq)
+      Util.writeArgs(dataType, inputPrepDataStream, arraySize, s: _*)
+
+    inputPrep.toByteArray()
+  }
+
+  def assertOutput(
+      dataType: ArchitectureDataType,
+      arraySize: Int,
+      bytes: Array[Byte],
+      count: Int = 1
+  ): Unit = {
+    val rmse = new RMSE()
+
+    val output =
+      new DataInputStream(new ByteArrayInputStream(bytes))
+
+    val outputSize = arraySize * 4
+
+    val result = Util
+      .readResult(dataType, output, arraySize, outputSize)
+      .toArray
+
+    val golden = Golden
+      .grouped(2)
+      .map(_ ++ Array.fill(arraySize - 2)(0f))
+      .flatten
+      .toArray
+
+    for (i <- 0 until outputSize)
+      rmse.addSample(result(i), golden(i))
+
+    assert(rmse.compute < dataType.error)
+  }
+
+  private val Golden = Seq(
+    1f, 5f, 2f, 6f, 3f, 7f, 4f, 8f
+  )
+}