docstr: Test multi device batch normalization at the operation level. This is
  done by checking the outputs of batch normalization and its gradient
  operator. We compare values produced with our manually calculated
  batch normalization values and gradients.
docstr_tok: ["Test","multi","device","batch","normalization","at","the","operation","level",".","This","is","done","by","checking","the","outputs","of","batch","normalization","and","its","gradient","operator",".","We","compare","values","produced","with","our","manually","calculated","batch","normalization","values","and","gradients","."]
code: def _bn_check_op_level(self, device_type, seed, batch_size):
      """
      Test multi device batch normalization at the operation level. This is
      done by checking the outputs of batch normalization and its gradient
      operator. We compare values produced with our manually calculated
      batch normalization values and gradients.
      """
      devices = [0, 1]
      epsilon = 1e-3
      tolerance = 1e-3
  
      def _test_forward_pass(
          x, devices, device_type, scale, bias, epsilon
      ):
          x_concat = np.concatenate(x)
          mean = np.mean(x_concat, axis=0)
          var = np.var(x_concat, axis=0)
          for device in devices:
              x_i = x[device]
              x_hat = (x_i - mean) / (np.sqrt(var + epsilon))
              expected_out = scale * x_hat + bias
              spatial_out = workspace.FetchBlob(
                  "{}_{}/bn_out".format(device_type, device)
              )
              rel_error = np.linalg.norm(
                  spatial_out - expected_out
              ) / np.linalg.norm(expected_out)
              self.assertTrue(rel_error < 0.005)
  
      def _test_backward_pass(
          x, devices, device_type, scale, tolerance
      ):
          dBias_arr = []
          dY_arr = []
          dGamma_arr = []
          num_devices = len(devices)
          mean = np.array(
              workspace.FetchBlob("{}_0/bn_out_sm".format(device_type)),
              dtype=np.float32,
          )
          inv_var = np.array(
              workspace.FetchBlob(
                  "{}_0/bn_out_siv".format(device_type)
              ),
              dtype=np.float32,
          )
  
          # dBias
          # Sum dBias values over all devices to find the average gradient
          for device in devices:
              dY_blob = workspace.FetchBlob(
                  "{}_{}/bn_out_grad".format(device_type, device)
              )
              dY = np.array(dY_blob, dtype=np.float32)
              dY_arr.append(dY)
              dBias_arr.append(
                  np.array(np.sum(dY, axis=0), dtype=np.float32)
              )
          dBias = np.sum(dBias_arr, dtype=np.float32)
          dBias_avg = dBias / num_devices
          for device in devices:
              dBiasActual = np.sum(
                  workspace.FetchBlob(
                      "{}_{}/bn_out_b_grad".format(device_type, device)
                  ),
                  dtype=np.float32,
              )
              self.assertTrue(
                  np.isclose([dBiasActual], [dBias], atol=tolerance)
              )
  
          # dGamma
          # Sum dGamma values over all devices to find the average gradient
          for device in devices:
              dGamma = np.sum(
                  (x[device] - mean) * inv_var * dY_arr[device],
                  axis=0,
                  dtype=np.float32,
              )
              dGamma_arr.append(dGamma)
          dGamma = np.sum(dGamma_arr, axis=0, dtype=np.float32)
          dGamma_avg = dGamma / num_devices
          for device in devices:
              dGammaActual = workspace.FetchBlob(
                  "{}_{}/bn_out_s_grad".format(device_type, device)
              )
              self.assertTrue(
                  np.isclose([dGamma], [dGammaActual], atol=tolerance)
              )
  
          # dX
          scale_inv_var = scale * inv_var / batch_size
          for device in devices:
              dX = scale_inv_var * (
                  dY_arr[device] * batch_size
                  - dBias_avg
                  - (x[device] - mean) * dGamma_avg * inv_var
              )
              dX_actual = workspace.FetchBlob(
                  "{}_{}/tanh_grad".format(device_type, device)
              )
              self.assertTrue(
                  np.isclose([dX], [dX_actual], atol=tolerance).all()
              )
  
      def add_input_ops(model):
          for device in devices:
              data = np.random.rand(batch_size, 1, 1, 1).astype(
                  np.float32
              )
              workspace.FeedBlob(
                  "{}_{}/data".format(device_type, device), data
              )
  
      def add_model_ops(model, loss_scale):
          if device_type == "gpu":
              model.CopyCPUToGPU("data", "device_data")
              model.Tanh("device_data", "tanh")
          else:
              model.Tanh("data", "tanh")
          model.SpatialBN(
              "tanh", "bn_out", 1, epsilon=epsilon, is_test=False
          )
          model.Sqr("bn_out", "sqr")
          loss = model.SumElements("sqr", "loss")
          return [loss]
  
      def add_optimizer(model):
          return optimizer.build_sgd(model, 0.1)
  
      np.random.seed(seed)
      workspace.ResetWorkspace()
      model = cnn.CNNModelHelper(order="NCHW", name="test")
      data_parallel_model.Parallelize(
          model,
          input_builder_fun=add_input_ops,
          forward_pass_builder_fun=add_model_ops,
          optimizer_builder_fun=add_optimizer,
          devices=devices,
          cpu_device=device_type == "cpu",
          shared_model=False,
          combine_spatial_bn=True,
      )
  
      workspace.RunNetOnce(model.param_init_net)
      scale = workspace.FetchBlob("{}_0/bn_out_s".format(device_type))
      bias = workspace.FetchBlob("{}_0/bn_out_b".format(device_type))
      workspace.RunNetOnce(model.net)
  
      x = []
      for device in devices:
          x_blob = workspace.FetchBlob(
              "{}_{}/tanh".format(device_type, device)
          )
          x_i = np.array(x_blob, dtype=np.float32)
          x.append(x_i)
  
      _test_forward_pass(x, devices, device_type, scale, bias, epsilon)
      _test_backward_pass(x, devices, device_type, scale, tolerance)
  
imports: from multiprocessing import Process, Queue
  import numpy
  import os
  import shutil
  import tempfile
  import unittest
  import time
  from unittest.mock import Mock
  from hypothesis import assume, given, settings
  import hypothesis.strategies
  from caffe2.proto import caffe2_pb2
  from caffe2.python import brew, core, cnn, data_parallel_model, dyndep, model_helper, optimizer, rnn_cell, workspace
  from caffe2.python.test_util import TestCase
label_desc: Use image node_id 12 for calling the DataParallelModelTest obj's underlying member method code with example usage: obj._bn_check_op_level(device_type, seed, batch_size) and returns: list, optimizer