docstr: Helper function for test_equiv
docstr_tok: ["Helper","function","for","test_equiv"]
code: def run_model(self, devices, gpu):
      """
      Helper function for test_equiv
      """
  
      def input_builder_fun(model):
          return None
  
      def model_build_fun(model, loss_scale):
          workspace.FeedBlob(
              core.ScopedBlobReference("seq_lengths"),
              np.array(
                  [self.T] * self.batch_per_device, dtype=np.int32
              ),
          )
          model.param_init_net.ConstantFill(
              [],
              "hidden_init",
              value=0.0,
              shape=[1, self.batch_per_device, self.hidden_dim],
          )
          model.param_init_net.ConstantFill(
              [],
              "cell_init",
              value=0.0,
              shape=[1, self.batch_per_device, self.hidden_dim],
          )
  
          (
              output,
              _last_hidden,
              _,
              _last_state,
          ) = rnn_cell.LSTM(
              model=model,
              input_blob="data",
              seq_lengths="seq_lengths",
              initial_states=("hidden_init", "cell_init"),
              dim_in=self.input_dim,
              dim_out=self.hidden_dim,
              scope="partest",
          )
  
          # A silly loss function
          loss = model.AveragedLoss(
              model.Sub([output, "target"], "dist"),
              "loss",
          )
          loss = model.Scale(loss, "loss_scaled", scale=loss_scale)
          return [loss]
  
      def param_update_fun(model):
          ITER = model.Iter("ITER")
          LR = model.net.LearningRate(
              [ITER],
              "LR",
              base_lr=(-0.1),
              policy="fixed",
          )
          ONE = model.param_init_net.ConstantFill(
              [],
              "ONE",
              shape=[1],
              value=1.0,
          )
          for param in model.GetParams():
              param_grad = model.param_to_grad[param]
              model.WeightedSum([param, ONE, param_grad, LR], param)
  
          assert len(model.GetParams()) == len(model.params) // len(
              model._devices
          )
  
      workspace.ResetWorkspace()
      model = cnn.CNNModelHelper(
          name="recurrent_test{}".format(devices),
      )
  
      self.T = 8
      self.batch_size = 64
      self.input_dim = 8
      self.hidden_dim = 31
      self.batch_per_device = self.batch_size // len(devices)
  
      data_parallel_model.Parallelize(
          model,
          input_builder_fun=input_builder_fun,
          forward_pass_builder_fun=model_build_fun,
          param_update_builder_fun=param_update_fun,
          devices=devices,
          optimize_gradient_memory=True,
          cpu_device=not gpu,
      )
  
      # Change all initialization to be ConstantFills so that
      # the everything is deterministic
      for op in model.param_init_net.Proto().op:
          if op.type.endswith("Fill"):
              op.type = "ConstantFill"
  
      # Each run has same input, independent of number of gpus
      np.random.seed(20150210)
      for i in range(0, 10):
          full_data = np.random.rand(
              self.T, self.batch_size, self.input_dim
          )
          full_target = np.random.rand(
              self.T, self.batch_size, self.hidden_dim
          )
  
          for j, g in enumerate(devices):
              st = j * self.batch_per_device
              en = st + self.batch_per_device
              data = full_data[:, st:en, :].astype(np.float32)
              targets = full_target[:, st:en, :].astype(np.float32)
              with core.DeviceScope(
                  core.DeviceOption(model._device_type, g)
              ):
                  workspace.FeedBlob(
                      "{}_{}/data".format(model._device_prefix, g), data
                  )
                  workspace.FeedBlob(
                      "{}_{}/target".format(model._device_prefix, g),
                      targets,
                  )
  
          if i == 0:
              workspace.RunNetOnce(model.param_init_net)
              workspace.CreateNet(model.net)
  
          workspace.RunNet(model.net.Proto().name)
  
      return workspace.FetchBlob(
          "{}_0/partest/i2h_w".format(model._device_prefix)
      )
  
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 1 for calling the RecurrentNetworkParallelTest obj's underlying member method code with example usage: obj.run_model(devices, gpu) and returns: workspace, None, list