Reverting inference behavior for BatchNorm. (#426)

BradLarson · web-flow · commit 268055eec408 · 2019-08-08T17:05:04.000-05:00
diff --git a/Sources/TensorFlow/Layers/Normalization.swift b/Sources/TensorFlow/Layers/Normalization.swift
@@ -82,10 +82,8 @@ public struct BatchNorm<Scalar: TensorFlowFloatingPoint>: Layer {
           let inv = rsqrt(moments.variance + epsilon) * scale.reshaped(to: moments.variance.shape)
           return (input - moments.mean) * inv + offset.reshaped(to: moments.mean.shape)
         case .inference:
-          let scaleShape = runningVariance.value.shape
-          let offsetShape = runningMean.value.shape
-          let inv = rsqrt(runningVariance.value + epsilon) * scale.reshaped(to: scaleShape)
-          return (input - runningMean.value) * inv + offset.reshaped(to: offsetShape)
+          let inv = rsqrt(runningVariance.value + epsilon) * scale
+          return (input - runningMean.value) * inv + offset
         }
     }
 

Original file line number	Diff line number	Diff line change
`@@ -82,10 +82,8 @@ public struct BatchNorm<Scalar: TensorFlowFloatingPoint>: Layer {`
`82`	`82`	`let inv = rsqrt(moments.variance + epsilon) * scale.reshaped(to: moments.variance.shape)`
`83`	`83`	`return (input - moments.mean) * inv + offset.reshaped(to: moments.mean.shape)`
`84`	`84`	`case .inference:`
`85`		`- let scaleShape = runningVariance.value.shape`
`86`		`- let offsetShape = runningMean.value.shape`
`87`		`- let inv = rsqrt(runningVariance.value + epsilon) * scale.reshaped(to: scaleShape)`
`88`		`- return (input - runningMean.value) * inv + offset.reshaped(to: offsetShape)`
	`85`	`+ let inv = rsqrt(runningVariance.value + epsilon) * scale`
	`86`	`+ return (input - runningMean.value) * inv + offset`
`89`	`87`	`}`
`90`	`88`	`}`
`91`	`89`