From 7a1a275472f9ebc0aa4212debd8051bf18e0dce2 Mon Sep 17 00:00:00 2001
From: Yang Bo <atryyang@thoughtworks.com>
Date: Thu, 31 May 2018 17:48:10 +0800
Subject: [PATCH] Add the excludeUnmatchedFineGrainedNetwork parameter

---
 .../deeplearning/benchmark/benchmark.scala    | 76 ++++++++++++++++---
 1 file changed, 66 insertions(+), 10 deletions(-)

diff --git a/benchmark/src/jmh/scala/com/thoughtworks/deeplearning/benchmark/benchmark.scala b/benchmark/src/jmh/scala/com/thoughtworks/deeplearning/benchmark/benchmark.scala
index 2875b5ec..4be73286 100644
--- a/benchmark/src/jmh/scala/com/thoughtworks/deeplearning/benchmark/benchmark.scala
+++ b/benchmark/src/jmh/scala/com/thoughtworks/deeplearning/benchmark/benchmark.scala
@@ -40,9 +40,9 @@ object benchmark {
 
   @Threads(value = 1)
   @State(Scope.Benchmark)
-  class FourLayer {
+  class BranchNetBenchmark {
 
-    @Param(Array("8"))
+    @Param(Array("8", "16"))
     protected var batchSize: Int = _
 
     @Param(Array("1", "2", "4"))
@@ -54,6 +54,9 @@ object benchmark {
     @Param(Array("16", "8"))
     protected var numberOfBranches: Int = _
 
+    @Param(Array("false", "true"))
+    protected var excludeUnmatchedFineGrainedNetwork: Boolean = _
+
     private implicit var executionContext: ExecutionContextExecutorService = _
 
     private lazy val batches = {
@@ -103,7 +106,7 @@ object benchmark {
         }
       }
 
-      val fineProbabilityModel = Seq.fill(Cifar100.NumberOfCoarseClasses)(new (INDArrayLayer => INDArrayLayer) {
+      val fineScoreModels = Seq.fill(Cifar100.NumberOfCoarseClasses)(new (INDArrayLayer => INDArrayLayer) {
         object Dense2 extends (INDArrayLayer => INDArrayLayer) {
 
           object Dense1 extends (INDArrayLayer => INDArrayLayer) {
@@ -127,14 +130,42 @@ object benchmark {
         val bias = INDArrayWeight(Nd4j.randn(1, Cifar100.NumberOfFineClassesPerCoarseClass))
 
         def apply(coarseFeatures: INDArrayLayer) = {
-          val scores = Dense2(coarseFeatures) dot weight + bias
-
-          val expScores = exp(scores)
-          expScores / expScores.sum(1)
+          Dense2(coarseFeatures) dot weight + bias
         }
       })
 
-      def loss(coarseLabel: Int, batch: Batch): DoubleLayer = {
+//      val fineProbabilityModel = Seq.fill(Cifar100.NumberOfCoarseClasses)(new (INDArrayLayer => INDArrayLayer) {
+//        object Dense2 extends (INDArrayLayer => INDArrayLayer) {
+//
+//          object Dense1 extends (INDArrayLayer => INDArrayLayer) {
+//            val weight = INDArrayWeight(Nd4j.randn(numberOfHiddenFeatures, numberOfHiddenFeatures))
+//            val bias = INDArrayWeight(Nd4j.randn(1, numberOfHiddenFeatures))
+//
+//            def apply(coarseFeatures: INDArrayLayer) = {
+//              max(coarseFeatures dot weight + bias, 0.0)
+//            }
+//          }
+//
+//          val weight = INDArrayWeight(Nd4j.randn(numberOfHiddenFeatures, numberOfHiddenFeatures))
+//          val bias = INDArrayWeight(Nd4j.randn(1, numberOfHiddenFeatures))
+//
+//          def apply(coarseFeatures: INDArrayLayer) = {
+//            max(Dense1(coarseFeatures) dot weight + bias, 0.0)
+//          }
+//        }
+//
+//        val weight = INDArrayWeight(Nd4j.randn(numberOfHiddenFeatures, Cifar100.NumberOfFineClassesPerCoarseClass))
+//        val bias = INDArrayWeight(Nd4j.randn(1, Cifar100.NumberOfFineClassesPerCoarseClass))
+//
+//        def apply(coarseFeatures: INDArrayLayer) = {
+//          val scores = Dense2(coarseFeatures) dot weight + bias
+//
+//          val expScores = exp(scores)
+//          expScores / expScores.sum(1)
+//        }
+//      })
+
+      def loss(expectedCoarseLabel: Int, batch: Batch): DoubleLayer = {
         def crossEntropy(prediction: INDArrayLayer, expectOutput: INDArray): DoubleLayer = {
           -(hyperparameters.log(prediction) * expectOutput).mean
         }
@@ -142,9 +173,34 @@ object benchmark {
         val Array(batchSize, width, height, channels) = batch.pixels.shape()
         val coarseFeatures = CoarseFeatures(batch.pixels.reshape(batchSize, width * height * channels))
         val coarseProbabilities = CoarseProbabilityModel(coarseFeatures)
-        val fineProbabilities = fineProbabilityModel(coarseLabel)(coarseFeatures)
 
-        crossEntropy(coarseProbabilities, batch.coarseClasses) + crossEntropy(fineProbabilities, batch.localFineClasses)
+        crossEntropy(coarseProbabilities, batch.coarseClasses) + {
+          if (excludeUnmatchedFineGrainedNetwork) {
+            val fineScores = fineScoreModels(expectedCoarseLabel)(coarseFeatures)
+            val expScores = exp(fineScores)
+            val fineProbabilities = expScores / expScores.sum(1)
+            crossEntropy(fineProbabilities, batch.localFineClasses)
+          } else {
+            val expScoresByCoarseLabel = for (coarseLabel <- 0 until Cifar100.NumberOfCoarseClasses) yield {
+              val fineScores = fineScoreModels(expectedCoarseLabel)(coarseFeatures)
+              exp(fineScores)
+            }
+            val expSum = expScoresByCoarseLabel.map(_.sum(1)).reduce(_ + _)
+            val lossPerCoarseLabel = for ((expScores, coarseLabel) <- expScoresByCoarseLabel.zipWithIndex) yield {
+              val fineProbabilities = expScores / expSum
+
+              crossEntropy(
+                fineProbabilities,
+                if (coarseLabel == expScoresByCoarseLabel) {
+                  batch.localFineClasses
+                } else {
+                  Nd4j.zeros(batchSize, Cifar100.NumberOfFineClassesPerCoarseClass)
+                }
+              )
+            }
+            lossPerCoarseLabel.reduce(_ + _)
+          }
+        }
       }
 
       def train(coarseLabel: Int, batch: Batch) = {