From ff0a7ae43f9b814801d7fb7144f4d390c586838b Mon Sep 17 00:00:00 2001
From: Jerry Zhang <jerryzh@fb.com>
Date: Tue, 16 Apr 2019 18:57:13 -0700
Subject: [PATCH] Testing for folded conv_bn_relu (#19298)

Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/19298

Proper testing for conv_bn_relu folding

Differential Revision: D13998891

fbshipit-source-id: ceb58ccec19885cbbf38964ee0d0db070e098b4a
---
 caffe2/python/brew.py                  |  1 +
 caffe2/python/helpers/normalization.py | 46 ++++++++++++++++++++++++++++++++++
 2 files changed, 47 insertions(+)

diff --git a/caffe2/python/brew.py b/caffe2/python/brew.py
index bb6d3f0..2722c21 100644
--- a/caffe2/python/brew.py
+++ b/caffe2/python/brew.py
@@ -45,6 +45,7 @@ class HelperWrapper(object):
         'instance_norm': instance_norm,
         'spatial_bn': spatial_bn,
         'spatial_gn': spatial_gn,
+        'moments_with_running_stats': moments_with_running_stats,
         'relu': relu,
         'prelu': prelu,
         'tanh': tanh,
diff --git a/caffe2/python/helpers/normalization.py b/caffe2/python/helpers/normalization.py
index a47ac58..6c0889c 100644
--- a/caffe2/python/helpers/normalization.py
+++ b/caffe2/python/helpers/normalization.py
@@ -290,3 +290,49 @@ def layer_norm(
     )
 
     return biased, mean, stdev
+
+def moments_with_running_stats(model, blob_in, blob_out, dim_in,
+                                     RunningMeanInitializer=None, RunningVarianceInitializer=None,
+                                     order="NCHW", **kwargs):
+
+    if model.init_params:
+        rm_init = ("ConstantFill", {'value': 0.0})
+        riv_init = ("ConstantFill", {'value': 1.0})
+
+        RunningMeanInitializer = initializers.update_initializer(
+            RunningMeanInitializer, rm_init, ("ConstantFill", {})
+        )
+        RunningVarianceInitializer = initializers.update_initializer(
+            RunningVarianceInitializer, riv_init, ("ConstantFill", {})
+        )
+    else:
+        RunningMeanInitializer = initializers.ExternalInitializer()
+        RunningVarianceInitializer = initializers.ExternalInitializer()
+
+    running_mean = model.create_param(
+        param_name=blob_out + '_rm',
+        shape=[dim_in],
+        initializer=RunningMeanInitializer,
+        tags=ParameterTags.COMPUTED_PARAM
+    )
+
+    # this is just running variance
+    running_inv_var = model.create_param(
+        param_name=blob_out + '_riv',
+        shape=[dim_in],
+        initializer=RunningVarianceInitializer,
+        tags=ParameterTags.COMPUTED_PARAM
+    )
+
+    blob_outs = [blob_out + "_sm", blob_out + "_sv"]
+    if order == 'NCHW':
+        blob_outputs = model.net.Moments(
+            [blob_in], blob_outs,
+            axes=[0, 2, 3],
+            order=order, keepdims=False, **kwargs)
+    elif order == 'NHWC':
+        blob_outputs = model.net.Moments(
+            [blob_in], blob_outs,
+            axes=[0, 1, 2],
+            order=order, keepdims=False, **kwargs)
+    return blob_outputs
-- 
2.7.4