Add OpStructureExporter.

morph_net/tools/json_tensor_exporter_op.cc

@@ -0,0 +1,23 @@
+#include "third_party/tensorflow/core/framework/op.h"
+
+
+REGISTER_OP("JsonTensorExporter")
+    .Attr("filename: string")
+    .Attr("T: {float, double, int32, bool}")
+    .Attr("N: int")
+    .Attr("keys: list(string)")
+    .Input("values: N * T")
+    .Input("save: bool")
+    .Doc(R"doc(
+Saves the content of tensors on file as JSON dictionary.
+
+filename: Filename to which the JSON is to be saved.
+N: Number of tensors expected.
+keys: The list of keys of the dictionary. Must be of length N.
+values: A list of tensors, will be the respective values. The order of the
+  values is expected to match that of the keys. Must be of length N. Currently
+  only vectors and scalars (rank 1 and 0) are supported.
+save: If false, the op would be a no-op. This mechanism is introduced because
+  tf.cond can execute both the if and the else, and we don't want to write files
+  unnecessarily.
+)doc");

morph_net/tools/json_tensor_exporter_op_kernel.cc

@@ -0,0 +1,108 @@
+#include <string>
+
+#include "file/base/file.h"
+#include "file/base/helpers.h"
+
+#include "file/base/options.h"
+#include "third_party/jsoncpp/json.h"
+#include "third_party/tensorflow/core/framework/op_kernel.h"
+#include "third_party/tensorflow/core/framework/tensor.h"
+#include "third_party/tensorflow/core/framework/tensor_shape.h"
+#include "third_party/tensorflow/core/lib/core/errors.h"
+
+
+namespace morph_net {
+
+using ::tensorflow::errors::InvalidArgument;
+using ::tensorflow::OpInputList;
+using ::tensorflow::OpKernelConstruction;
+using ::tensorflow::OpKernel;
+using ::tensorflow::OpKernelContext;
+using ::tensorflow::Tensor;
+
+
+template <typename T>
+class JsonTensorExporterOpKernel : public OpKernel {
+ public:
+  explicit JsonTensorExporterOpKernel(OpKernelConstruction* context)
+      : OpKernel(context) {
+    int number_of_keys;
+    OP_REQUIRES_OK(context, context->GetAttr("N", &number_of_keys));
+    OP_REQUIRES_OK(context, context->GetAttr("keys", &keys_));
+    OP_REQUIRES_OK(context, context->GetAttr("filename", &filename_));
+
+    OP_REQUIRES(context, keys_.size() == number_of_keys,
+                InvalidArgument("Number of keys (", keys_.size(), ") must match"
+                                " N (", number_of_keys, ")."));
+
+    OP_REQUIRES_OK(context, WriteFile(""));
+  }
+
+  void Compute(OpKernelContext* context) override {
+    OpInputList values;
+    const Tensor* save;
+    OP_REQUIRES_OK(context, context->input_list("values", &values));
+    OP_REQUIRES_OK(context, context->input("save", &save));
+    if (!save->scalar<bool>()()) return;
+
+    CHECK_EQ(values.size(), keys_.size());  // Enforced by REGISTER_OP
+
+    Json::Value json;
+    int ikey = 0;
+    for (const Tensor& tensor : values) {
+       OP_REQUIRES(context, tensor.dims() <= 1, InvalidArgument(
+           "Only scalars and vectors are currnetly supported, but a tensor "
+           "with rank ", tensor.dims(), "was found."));
+
+      const string& key = keys_[ikey++];
+      if (tensor.dims() == 0) {  // Scalar
+        json[key] = tensor.scalar<T>()();
+        continue;
+      }
+
+      // Vector
+      for (int ielement = 0; ielement < tensor.NumElements(); ++ielement) {
+        json[key][ielement] = tensor.vec<T>()(ielement);
+      }
+    }
+
+    Json::StyledWriter writer;
+    OP_REQUIRES_OK(context, WriteFile(writer.write(json)));
+  }
+
+ private:
+  ::tensorflow::Status WriteFile(const string& content) const {
+    ::util::Status status =
+        ::file::SetContents(filename_, content, ::file::Defaults());
+    if (status.ok()){
+      return ::tensorflow::Status::OK();
+    }
+    return InvalidArgument("Unable to write to file ", filename_,
+                           ". Error message: ", status.error_message());
+  }
+
+  std::vector<string> keys_;
+  string filename_;
+};
+
+REGISTER_KERNEL_BUILDER(Name("JsonTensorExporter")
+                            .Device(::tensorflow::DEVICE_CPU)
+                            .TypeConstraint<int32>("T"),
+                        JsonTensorExporterOpKernel<int32>);
+
+REGISTER_KERNEL_BUILDER(Name("JsonTensorExporter")
+                            .Device(::tensorflow::DEVICE_CPU)
+                            .TypeConstraint<float>("T"),
+                        JsonTensorExporterOpKernel<float>);
+
+REGISTER_KERNEL_BUILDER(Name("JsonTensorExporter")
+                            .Device(::tensorflow::DEVICE_CPU)
+                            .TypeConstraint<double>("T"),
+                        JsonTensorExporterOpKernel<double>);
+
+REGISTER_KERNEL_BUILDER(Name("JsonTensorExporter")
+                            .Device(::tensorflow::DEVICE_CPU)
+                            .TypeConstraint<bool>("T"),
+                        JsonTensorExporterOpKernel<bool>);
+
+}  // namespace morph_net

morph_net/tools/json_tensor_exporter_op_test.cc

@@ -0,0 +1,118 @@
+#include <string>
+
+#include "file/base/file.h"
+#include "file/base/helpers.h"
+#include "file/base/path.h"
+
+#include "testing/base/public/gmock.h"
+#include "testing/base/public/gunit.h"
+#include "third_party/jsoncpp/json.h"
+#include "third_party/tensorflow/core/framework/fake_input.h"
+#include "third_party/tensorflow/core/framework/node_def_builder.h"
+#include "third_party/tensorflow/core/framework/tensor.h"
+#include "third_party/tensorflow/core/framework/tensor_testutil.h"
+#include "third_party/tensorflow/core/kernels/ops_testutil.h"
+#include "third_party/tensorflow/core/lib/core/status_test_util.h"
+
+namespace morph_net {
+
+using ::tensorflow::DT_INT32;
+using ::tensorflow::FakeInput;
+using ::tensorflow::NodeDefBuilder;
+using ::tensorflow::OpsTestBase;
+using ::tensorflow::Status;
+using ::tensorflow::TensorShape;
+using ::testing::ElementsAre;
+
+
+std::vector<int> ToVector(const Json::Value& json) {
+  std::vector<int> v;
+  for (const Json::Value& item : json) {
+    v.push_back(item.asInt());
+  }
+  return v;
+}
+
+
+class JsonTensorExporterTest : public OpsTestBase {};
+
+TEST_F(JsonTensorExporterTest, Success) {
+  const int kLength = 3;
+  const string filename = ::file::JoinPath(FLAGS_test_tmpdir, "success.json");
+  TF_ASSERT_OK(
+       NodeDefBuilder("exporter", "JsonTensorExporter")
+      .Attr("T", DT_INT32)
+      .Attr("N", kLength)
+      .Attr("keys", {"k1", "k2", "k3"})
+      .Attr("filename", filename)
+      .Input(FakeInput(kLength, ::tensorflow::DT_INT32))
+      .Input(FakeInput(::tensorflow::DT_BOOL))
+      .Finalize(node_def()));
+
+  TF_ASSERT_OK(InitOp());
+  // The initialization of the op creates an empty file at `filename`. We delete
+  // both to verify it was created, and to clean it up for the next steps of the
+  // test.
+  ASSERT_OK(::file::Delete(filename, ::file::Defaults()));
+
+  AddInputFromArray<int>(TensorShape({3}), {3, 5, 7});
+  AddInputFromArray<int>(TensorShape({2}), {6, 4});
+  AddInputFromArray<int>(TensorShape({}), {9});
+
+  // Set the `save` flag initially to false - so the op should be a no-op.
+  AddInputFromArray<bool>(TensorShape({}), {false});
+  TF_ASSERT_OK(RunOpKernel());
+  // Verify that indeed no file was created.
+  EXPECT_EQ(absl::StatusCode::kNotFound,
+            ::file::Exists(filename, ::file::Defaults()).code());
+
+  // Flip the `save` flag to true and test the content of the savef file.
+  tensors_[3]->scalar<bool>()() = true;
+  TF_ASSERT_OK(RunOpKernel());
+
+  string contents;
+  ASSERT_OK(::file::GetContents(filename, &contents, ::file::Defaults()));
+  Json::Reader reader;
+  Json::Value json;
+  reader.parse(contents, json);
+  EXPECT_THAT(json.getMemberNames(), ElementsAre("k1", "k2", "k3"));
+  EXPECT_TRUE(json["k1"].isArray());
+  EXPECT_THAT(ToVector(json["k1"]), ElementsAre(3, 5, 7));
+  EXPECT_TRUE(json["k2"].isArray());
+  EXPECT_THAT(ToVector(json["k2"]), ElementsAre(6, 4));
+  EXPECT_EQ(9, json["k3"].asInt());
+}
+
+TEST_F(JsonTensorExporterTest, WrongNumberOfKeys) {
+  const int kLength = 3;
+  const string filename = ::file::JoinPath(FLAGS_test_tmpdir, "failure.json");
+  TF_ASSERT_OK(
+       NodeDefBuilder("exporter", "JsonTensorExporter")
+      .Attr("T", DT_INT32)
+      .Attr("N", kLength)
+      .Attr("keys", {"k1", "k2"})  // Two keys only, even though kLength = 3.
+      .Attr("filename", filename)
+      .Input(FakeInput(kLength, ::tensorflow::DT_INT32))
+      .Input(FakeInput(::tensorflow::DT_BOOL))
+      .Finalize(node_def()));
+
+  EXPECT_FALSE(InitOp().ok());
+}
+
+TEST_F(JsonTensorExporterTest, BadFileName) {
+  const int kLength = 3;
+  const string filename = "**bad";
+  TF_ASSERT_OK(
+       NodeDefBuilder("exporter", "JsonTensorExporter")
+      .Attr("T", DT_INT32)
+      .Attr("N", kLength)
+      .Attr("keys", {"k1", "k2", "k3"})
+      .Attr("filename", filename)
+      .Input(FakeInput(kLength, ::tensorflow::DT_INT32))
+      .Input(FakeInput(::tensorflow::DT_BOOL))
+      .Finalize(node_def()));
+
+  EXPECT_FALSE(InitOp().ok());
+}
+
+}  // namespace morph_net