From 331c4b5b4d600049fbd5a6503c9cf5f1009214c5 Mon Sep 17 00:00:00 2001
From: Duc Ngo <duc@fb.com>
Date: Thu, 13 Dec 2018 20:43:00 -0800
Subject: [PATCH] caffe2 - make DataRandomFiller usable in unit tests (#15027)

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

- Make DataRandomFiller able to accept input_dims and input_types for only non intermediate inputs. Add a helper to fill input directly to a workspace

Reviewed By: highker

Differential Revision: D13408345

fbshipit-source-id: 5fc54d33da12e3f0a200e79380d4c695b0339b17
---
 caffe2/predictor/emulator/data_filler.cc | 86 ++++++++++++++++++++++++++++++++
 caffe2/predictor/emulator/data_filler.h  | 20 +++++++-
 2 files changed, 105 insertions(+), 1 deletion(-)

diff --git a/caffe2/predictor/emulator/data_filler.cc b/caffe2/predictor/emulator/data_filler.cc
index a0e6c6e..e4e64a3 100644
--- a/caffe2/predictor/emulator/data_filler.cc
+++ b/caffe2/predictor/emulator/data_filler.cc
@@ -159,5 +159,91 @@ void DataRandomFiller::fill_input_internal(TensorList_t* input_data) const {
   }
 }
 
+TestDataRandomFiller::TestDataRandomFiller(
+    const NetDef& net,
+    const std::vector<std::vector<std::vector<int64_t>>>& inputDims,
+    const std::vector<std::vector<std::string>>& inputTypes)
+    : DataRandomFiller() {
+  std::unordered_set<std::string> outputNames;
+  // Determine blobs that are outputs of some ops (intermediate blobs).
+  for (auto opIdx = 0; opIdx < net.op_size(); ++opIdx) {
+    const auto& op = net.op(opIdx);
+    for (auto outputIdx = 0; outputIdx < op.output_size(); ++outputIdx) {
+      outputNames.emplace(op.output(outputIdx));
+    }
+  }
+  // Determine ops that have non-intermediate inputs.
+  std::unordered_set<size_t> opWithRequiredInputs;
+  for (auto opIdx = 0; opIdx < net.op_size(); ++opIdx) {
+    const auto& op = net.op(opIdx);
+    for (auto inputIdx = 0; inputIdx < op.input_size(); ++inputIdx) {
+      if (!outputNames.count(op.input(inputIdx))) {
+        opWithRequiredInputs.emplace(opIdx);
+        break;
+      }
+    }
+  }
+
+  CAFFE_ENFORCE_EQ(inputDims.size(), opWithRequiredInputs.size());
+  CAFFE_ENFORCE_EQ(inputTypes.size(), opWithRequiredInputs.size());
+
+  int counter = 0;
+  for (auto opIdx = 0; opIdx < net.op_size(); ++opIdx) {
+    if (!opWithRequiredInputs.count(opIdx)) {
+      // Skip intermediate ops.
+      continue;
+    }
+    const auto& op = net.op(opIdx);
+    const auto& op_dims = inputDims[counter];
+    const auto& op_types = inputTypes[counter];
+    ++counter;
+
+    int countRequiredInputs = 0;
+    for (auto inputIdx = 0; inputIdx < op.input_size(); ++inputIdx) {
+      if (!outputNames.count(op.input(inputIdx))) {
+        ++countRequiredInputs;
+      }
+    }
+
+    CAFFE_ENFORCE(
+        op_dims.size() == countRequiredInputs,
+        op.name() + " has " + c10::to_string(op.input_size()) +
+            " (required) inputs; while the input dimension size is " +
+            c10::to_string(op_dims.size()));
+    CAFFE_ENFORCE(
+        op_types.size() == countRequiredInputs,
+        op.name() + " has " + c10::to_string(op.input_size()) +
+            " (required) inputs; while the input type size is " +
+            c10::to_string(op_types.size()));
+
+    int dimCounter = 0;
+    for (auto inputIdx = 0; inputIdx < op.input_size(); ++inputIdx) {
+      auto inputName = op.input(inputIdx);
+      if (outputNames.count(inputName)) {
+        // Skip intermediate inputs.
+        continue;
+      }
+      inputs_[inputName] = std::make_pair(
+          get_tensor_filler(op, dimCounter, op_dims), op_types[dimCounter]);
+      ++dimCounter;
+    }
+  }
+  CAFFE_ENFORCE(inputs_.size() > 0, "Empty input for run net");
+  // generate input names
+  for (const auto& input : inputs_) {
+    input_names_.push_back(input.first);
+  }
+}
+
+void TestDataRandomFiller::fillInputToWorkspace(Workspace* workspace) const {
+  for (auto& name : input_names_) {
+    const auto& it = inputs_.find(name);
+    CAFFE_ENFORCE(it != inputs_.end());
+    auto* tensor =
+        BlobGetMutableTensor(workspace->CreateBlob(name), caffe2::CPU);
+    fill_with_type(it->second.first, it->second.second, tensor);
+  }
+}
+
 } // namespace emulator
 } // namespace caffe2
diff --git a/caffe2/predictor/emulator/data_filler.h b/caffe2/predictor/emulator/data_filler.h
index 4c3d6a8..e574ba5 100644
--- a/caffe2/predictor/emulator/data_filler.h
+++ b/caffe2/predictor/emulator/data_filler.h
@@ -93,7 +93,9 @@ class DataRandomFiller : public Filler {
 
   void fill_parameter(Workspace* ws) const override;
 
- private:
+ protected:
+  DataRandomFiller() {}
+
   TensorFiller get_tensor_filler(
       const OperatorDef& op_def,
       int input_index,
@@ -118,5 +120,21 @@ class DataRandomFiller : public Filler {
   std::unordered_map<std::string, filler_type_pair_t> inputs_;
 };
 
+// A DataRandomFiller that is more convenient to use in unit tests.
+// Callers just need to supply input dimensions and types for non-intermediate
+// blobs.
+// It also treats parameters the same way as non-intermediate inputs (no
+// handling of parameters separately).
+class TestDataRandomFiller : public DataRandomFiller {
+ public:
+  TestDataRandomFiller(
+      const NetDef& net,
+      const std::vector<std::vector<std::vector<int64_t>>>& inputDims,
+      const std::vector<std::vector<std::string>>& inputTypes);
+
+  // Fill input directly to the workspace.
+  void fillInputToWorkspace(Workspace* workspace) const;
+};
+
 } // namespace emulator
 } // namespace caffe2
-- 
2.7.4