From 50da5bd7a1e61f8031c12f902ab67c5ef9177d55 Mon Sep 17 00:00:00 2001
From: Flowingsun007 <flowingsun007@163.com>
Date: Mon, 6 Jul 2020 11:28:24 +0800
Subject: [PATCH] add bn axis 3

---
 Classification/resnet50v1.5/resnet_model.py | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/Classification/resnet50v1.5/resnet_model.py b/Classification/resnet50v1.5/resnet_model.py
index 036a540..036e373 100755
--- a/Classification/resnet50v1.5/resnet_model.py
+++ b/Classification/resnet50v1.5/resnet_model.py
@@ -47,9 +47,12 @@ class ResnetBuilder(object):
 
     def _batch_norm(self, inputs, name=None, last=False):
         initializer = flow.zeros_initializer() if last else flow.ones_initializer()
+        axis = 1
+        if self.data_format =="NHWC":
+            axis = 3
         return flow.layers.batch_normalization(
             inputs=inputs,
-            axis=1,
+            axis=axis,
             momentum=0.9,  # 97,
             epsilon=1e-5,
             center=True,
-- 
GitLab