[MLGO] Use binary protobufs for improved training performance.

It turns out that during training, the time required to parse the
textual protobuf of a training log is about the same as the time it
takes to compile the module generating that log. Using binary protobufs
instead elides that cost almost completely.

Differential Revision: https://reviews.llvm.org/D106157

CMakeLists.txt

@@ -788,8 +788,26 @@ endif()
 set(TENSORFLOW_C_LIB_PATH "" CACHE PATH "Path to TensorFlow C library install")
 if (TENSORFLOW_C_LIB_PATH)
   find_library(tensorflow_c_api tensorflow PATHS ${TENSORFLOW_C_LIB_PATH}/lib NO_DEFAULT_PATH REQUIRED)
+  # Currently, the protobuf headers are distributed with the pip package that corresponds to the version
+  # of the C API library.
+  find_library(tensorflow_fx tensorflow_framework PATHS ${TENSORFLOW_C_LIB_PATH}/lib NO_DEFAULT_PATH REQUIRED)
   set(LLVM_HAVE_TF_API "ON" CACHE BOOL "Full Tensorflow API available")
   include_directories(${TENSORFLOW_C_LIB_PATH}/include)
+  if (NOT TF_PROTO_HEADERS)
+    message(STATUS "TF_PROTO_HEADERS not defined. Looking for tensorflow pip package.")
+    execute_process(COMMAND 
+      ${Python3_EXECUTABLE} "-m" "pip" "show" "tensorflow" 
+      OUTPUT_VARIABLE TF_PIP_OUT)
+    if ("${TF_PIP_OUT}" STREQUAL "")
+      message(FATAL ERROR "Tensorflow pip package is also required for 'development' mode (protobuf headers)")
+    endif()
+    string(REGEX MATCH "Location: ([^\n]*\n)" TF_PIP_LOC "${TF_PIP_OUT}")
+    string(REPLACE "Location: " "" TF_PIP ${TF_PIP_LOC})
+    set(TF_PROTO_HEADERS ${TF_PIP}/include)
+  endif()
+  include_directories(${TF_PROTO_HEADERS})
+  add_definitions("-DGOOGLE_PROTOBUF_NO_RTTI")
+  add_definitions("-D_GLIBCXX_USE_CXX11_ABI=0")
 endif()
 
 # For up-to-date instructions for installing the Tensorflow dependency, refer to

lib/Analysis/CMakeLists.txt

@@ -27,7 +27,7 @@ if (DEFINED LLVM_HAVE_TF_AOT OR DEFINED LLVM_HAVE_TF_API)
   endif()
 
   if (DEFINED LLVM_HAVE_TF_API)
-    list(APPEND MLLinkDeps ${tensorflow_c_api})
+    list(APPEND MLLinkDeps ${tensorflow_c_api} ${tensorflow_fx})
   endif()
 endif()
 

lib/Analysis/TFUtils.cpp

@@ -15,6 +15,7 @@
 
 #include "llvm/ADT/Twine.h"
 #include "llvm/Analysis/Utils/TFUtils.h"
+#include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/JSON.h"
 #include "llvm/Support/ManagedStatic.h"
@@ -22,14 +23,19 @@
 #include "llvm/Support/Path.h"
 #include "llvm/Support/raw_ostream.h"
 
+#include "google/protobuf/text_format.h"
 #include "tensorflow/c/c_api.h"
 #include "tensorflow/c/c_api_experimental.h"
-
+#include "tensorflow/core/example/example.pb.h"
 #include <cassert>
 #include <numeric>
 
 using namespace llvm;
 
+static cl::opt<bool>
+    ProtobufTextMode("tfutils-text-log", cl::init(false), cl::Hidden,
+                     cl::desc("Output textual (human-readable) protobuf."));
+
 namespace {
 
 using TFGraphPtr = std::unique_ptr<TF_Graph, decltype(&TF_DeleteGraph)>;
@@ -65,85 +71,53 @@ TFSessionOptionsPtr createTFSessionOptions() {
   return TFSessionOptionsPtr(TF_NewSessionOptions(), &TF_DeleteSessionOptions);
 }
 
-/// Write the values of one tensor as a list.
-template <typename T>
-void writeTensorValues(raw_ostream &OutFile, const char *TensorData,
-                       size_t ElemCount) {
-  OutFile << "[";
-  const T *TypedData = reinterpret_cast<const T *>(TensorData);
-  ListSeparator LS;
-  for (size_t I = 0; I < ElemCount; ++I)
-    OutFile << LS << TypedData[I];
-  OutFile << "]";
-}
-
 /// Write a list of tensors as a sequence of TensorFlow FeatureList protobufs.
 /// The tensors are assumed to be stored contiguously, in row-major format,
 /// in the TensorData buffer. Each tensor has the shape given by Spec. The
 /// feature name in the output is either the provided LoggingName, if
 /// specified, otherwise it's the name of the tensor (as given by Spec).
-void writeRawTensorsAsFeatureLists(raw_ostream &OutFile,
+void writeRawTensorsAsFeatureLists(tensorflow::FeatureLists *FE,
                                    const LoggedFeatureSpec &LoggedSpec,
                                    const char *TensorData, size_t TensorCount,
                                    bool FinalReward = false) {
-  const char *FieldName = "<invalid>";
-  std::function<void(const char *)> ValueWriter;
   const auto &Spec = LoggedSpec.Spec;
   // The 'Feature' protobuf only has 3 possible fields: float_list,
   // int64_list, or bytes_list, so we capture int32 values as int64. We don't
   // support any other types.
-  if (Spec.isElementType<int64_t>()) {
+  tensorflow::FeatureList &FL = (*FE->mutable_feature_list())[(
-    FieldName = "int64_list";
+      LoggedSpec.LoggingName ? *LoggedSpec.LoggingName : Spec.name())];
-    ValueWriter = [&](const char *Data) {
-      writeTensorValues<int64_t>(OutFile, Data, Spec.getElementCount());
-    };
-  } else if (Spec.isElementType<int32_t>()) {
-    FieldName = "int64_list";
-    ValueWriter = [&](const char *Data) {
-      writeTensorValues<int32_t>(OutFile, Data, Spec.getElementCount());
-    };
-
-  } else if (Spec.isElementType<float>()) {
-    FieldName = "float_list";
-    ValueWriter = [&](const char *Data) {
-      writeTensorValues<float>(OutFile, Data, Spec.getElementCount());
-    };
-
-  } else {
-    llvm_unreachable("Unsupported tensor type.");
-  }
-
-  OutFile << "  feature_list: {\n";
-  OutFile << "    key: "
-          << "\""
-          << (LoggedSpec.LoggingName ? *LoggedSpec.LoggingName : Spec.name())
-          << "\" ";
-  OutFile << "value: {\n";
-  size_t TensorByteSize = Spec.getElementCount() * Spec.getElementByteSize();
-
-  auto WriteFeatureProto = [&](const char *P) {
-    OutFile << "      feature: { " << FieldName << ": { value: ";
-    ValueWriter(P);
-    OutFile << " } }\n";
-  };
 
   const char *CurrentTensor = TensorData;
-  static int64_t Zero = 0;
+  const size_t TensorByteSize =
-  // Write all but the last value. If this is the final reward, don't increment
+      Spec.getElementCount() * Spec.getElementByteSize();
-  // the CurrentTensor, and just write 0.
+  const size_t ElemCount = Spec.getElementCount();
-  for (size_t I = 0; I < TensorCount - 1; ++I) {
+  for (size_t E = 0; E < TensorCount; ++E) {
-    if (FinalReward)
+    const bool ShouldWrite = E + 1 == TensorCount || !FinalReward;
-      WriteFeatureProto(reinterpret_cast<const char *>(&Zero));
+
-    else {
+    if (Spec.isElementType<int64_t>()) {
-      WriteFeatureProto(CurrentTensor);
+      auto *MF = FL.add_feature()->mutable_int64_list()->mutable_value();
-      CurrentTensor += TensorByteSize;
+      MF->Resize(ElemCount, 0);
+      if (ShouldWrite)
+        memcpy(MF->mutable_data(), CurrentTensor, TensorByteSize);
+    } else if (Spec.isElementType<int32_t>()) {
+      auto *MF = FL.add_feature()->mutable_int64_list()->mutable_value();
+      MF->Resize(ElemCount, 0);
+      if (ShouldWrite) {
+        const int32_t *TD = reinterpret_cast<const int32_t *>(CurrentTensor);
+        for (size_t I = 0; I < ElemCount; ++I)
+          (*MF)[I] = TD[I];
+      }
+    } else if (Spec.isElementType<float>()) {
+      auto *MF = FL.add_feature()->mutable_float_list()->mutable_value();
+      MF->Resize(ElemCount, 0.0);
+      if (ShouldWrite)
+        memcpy(MF->mutable_data(), CurrentTensor, TensorByteSize);
+    } else {
+      llvm_unreachable("Unsupported tensor type.");
     }
+    if (ShouldWrite)
+      CurrentTensor += TensorByteSize;
   }
-
-  WriteFeatureProto(CurrentTensor);
-
-  OutFile << "    }\n";
-  OutFile << "  }\n";
 }
 } // namespace
 
@@ -475,6 +449,8 @@ TFModelEvaluator::EvaluationResult::~EvaluationResult() {}
 TFModelEvaluator::~TFModelEvaluator() {}
 
 void Logger::print(raw_ostream &OS) {
+  tensorflow::SequenceExample SE;
+
   if (RawLogData.empty())
     return;
   if (RawLogData[0].empty())
@@ -488,16 +464,21 @@ void Logger::print(raw_ostream &OS) {
       RewardSpec.getElementCount() * RewardSpec.getElementByteSize();
   size_t NumberOfRewards = RawLogData.back().size() / RewardSize;
 
-  OS << "feature_lists: {\n";
+  tensorflow::FeatureLists *FE = SE.mutable_feature_lists();
   for (size_t I = 0; I < FeatureSpecs.size(); ++I)
-    writeRawTensorsAsFeatureLists(OS, FeatureSpecs[I], RawLogData[I].data(),
+    writeRawTensorsAsFeatureLists(FE, FeatureSpecs[I], RawLogData[I].data(),
                                   NumberOfRecords);
 
   if (IncludeReward)
-    writeRawTensorsAsFeatureLists(OS, {RewardSpec, None},
+    writeRawTensorsAsFeatureLists(FE, {RewardSpec, None},
                                   RawLogData.back().data(), NumberOfRecords,
                                   NumberOfRewards == 1);
-
+  std::string OutStr;
-  OS << "}\n";
+  if (ProtobufTextMode) {
+    google::protobuf::TextFormat::PrintToString(SE, &OutStr);
+  } else {
+    OutStr = SE.SerializeAsString();
+  }
+  OS << OutStr;
 }
 #endif // defined(LLVM_HAVE_TF_API)

test/Transforms/Inline/ML/bounds-checks-rewards.ll

@@ -7,40 +7,35 @@
 ; REQUIRES: have_tf_api
 ;
 ; Generate mock model
-; RUN: rm -rf %t && mkdir %t
+; RUN: rm -rf %t
 ; RUN: %python %S/../../../../lib/Analysis/models/generate_mock_model.py %S/../../../../lib/Analysis/models/inlining/config.py %t
 ;
 ; When the bounds are very wide ("no bounds"), all inlinings happen.
-; RUN: opt -passes=scc-oz-module-inliner -ml-inliner-ir2native-model=%S/../../../../unittests/Analysis/Inputs/ir2native_x86_64_model -ml-inliner-model-under-training=%t -training-log=- -enable-ml-inliner=development -ml-advisor-size-increase-threshold=10.0 -S < %s 2>&1 | FileCheck %s --check-prefix=CHECK --check-prefix=NOBOUNDS
+; RUN: opt -passes=scc-oz-module-inliner -ml-inliner-ir2native-model=%S/../../../../unittests/Analysis/Inputs/ir2native_x86_64_model -ml-inliner-model-under-training=%t -training-log=- -tfutils-text-log -enable-ml-inliner=development -ml-advisor-size-increase-threshold=10.0 -S < %s 2>&1 | FileCheck %s --check-prefix=CHECK --check-prefix=NOBOUNDS
 ;
 ; When the bounds are very restrictive, the first inlining happens but it's
 ; considered as "bad" (since it trips over the bounds) and its reward is a
 ; penalty. However, the mandatory inlining, which is considered next, happens.
 ; No other inlinings happend.
-; RUN: opt -passes=scc-oz-module-inliner -ml-inliner-ir2native-model=%S/../../../../unittests/Analysis/Inputs/ir2native_x86_64_model -ml-inliner-model-under-training=%t -training-log=- -enable-ml-inliner=development -ml-advisor-size-increase-threshold=1.0 -S < %s 2>&1 | FileCheck %s --check-prefix=CHECK --check-prefix=BOUNDS
+; RUN: opt -passes=scc-oz-module-inliner -ml-inliner-ir2native-model=%S/../../../../unittests/Analysis/Inputs/ir2native_x86_64_model -ml-inliner-model-under-training=%t -training-log=- -tfutils-text-log -enable-ml-inliner=development -ml-advisor-size-increase-threshold=1.0 -S < %s 2>&1 | FileCheck %s --check-prefix=CHECK --check-prefix=BOUNDS
 ;
 ; With more restrictive bounds, the first inlining happens and is OK. The
 ; mandatory inlining happens next, and it trips over the bounds, which then
 ; forces no further inlinings.
-; RUN: opt -passes=scc-oz-module-inliner -ml-inliner-ir2native-model=%S/../../../../unittests/Analysis/Inputs/ir2native_x86_64_model -ml-inliner-model-under-training=%t -training-log=- -enable-ml-inliner=development -ml-advisor-size-increase-threshold=1.1 -S < %s 2>&1 | FileCheck %s --check-prefix=CHECK --check-prefix=RELAXED-BOUNDS
+; RUN: opt -passes=scc-oz-module-inliner -ml-inliner-ir2native-model=%S/../../../../unittests/Analysis/Inputs/ir2native_x86_64_model -ml-inliner-model-under-training=%t -training-log=- -tfutils-text-log -enable-ml-inliner=development -ml-advisor-size-increase-threshold=1.1 -S < %s 2>&1 | FileCheck %s --check-prefix=CHECK --check-prefix=RELAXED-BOUNDS
-
 target datalayout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-f80:128-n8:16:32:64-S128"
 target triple = "x86_64-grtev4-linux-gnu"
-
 declare i64 @f1()
-
 define i64 @may_not_be_inlined() {
   %r = call i64 @f1()
   %r2 = add i64 13, %r
   ret i64 %r2
 }
-
 define i64 @must_be_inlined() #0 {
   %r = call i64 @may_not_be_inlined()
   %r2 = add i64 13, %r
   ret i64 %r2
 }
-
 define i64 @top() {
   %r = call i64 @must_be_inlined()
   %r2 = call i64 @may_not_be_inlined()
@@ -49,15 +44,25 @@ define i64 @top() {
   %r5 = add i64 %r3, %r4
   ret i64 %r5
 }
-
 attributes #0 = { alwaysinline }
-; CHECK: key: "delta_size" value: {
+; CHECK:       key: "delta_size"
-; NOBOUNDS-NEXT: feature: { int64_list: { value: [6] } }
+; CHECK-NEXT:     value {
-; RELAXED-BOUNDS-NEXT: feature: { int64_list: { value: [6] } }
+; CHECK-NEXT:       feature {
-; NOBOUNDS-NEXT: feature: { int64_list: { value: [-11] } }
+; CHECK-NEXT:         int64_list {
-; NOBOUNDS-NEXT: feature: { int64_list: { value: [4] } }
+; NOBOUNDS-NEXT:        value: 6
-; BOUNDS-NEXT: feature: { int64_list: { value: [2147483647] } }
+; RELAXED-BOUNDS-NEXT:  value: 6
-; CHECK-NEXT: }
+; NOBOUNDS-NEXT:      }
+; NOBOUNDS-NEXT:    }
+; NOBOUNDS-NEXT:    feature {
+; NOBOUNDS-NEXT:      int64_list {
+; NOBOUNDS-NEXT:        value: -11
+; NOBOUNDS-NEXT:      }
+; NOBOUNDS-NEXT:    }
+; NOBOUNDS-NEXT:    feature {
+; NOBOUNDS-NEXT:      int64_list {
+; NOBOUNDS-NEXT:        value: 4
+; BOUNDS-NEXT:          value: 2147483647
+; CHECK-NEXT:         }
 ; CHECK-LABEL: @top
 ; must_be_inlined must always be inlined, so we won't find a call to it in @top()
 ; CHECK-NOT: call i64 @must_be_inlined

test/Transforms/Inline/ML/development-training-log.ll

@@ -1,58 +1,56 @@
 ; Test that we can produce a log if we have or do not have a model, in development mode.
 ; REQUIRES: have_tf_api
 ; Generate mock model
-; RUN: rm -rf %t && mkdir %t
+; RUN: rm -rf %t
 ; RUN: %python %S/../../../../lib/Analysis/models/generate_mock_model.py %S/../../../../lib/Analysis/models/inlining/config.py %t
 ;
-; RUN: opt -enable-ml-inliner=development -passes=scc-oz-module-inliner -training-log=- -ml-inliner-model-under-training=%t -ml-inliner-ir2native-model=%S/../../../../unittests/Analysis/Inputs/ir2native_x86_64_model -S < %s | FileCheck %s 
+; RUN: opt -enable-ml-inliner=development -passes=scc-oz-module-inliner -training-log=- -tfutils-text-log -ml-inliner-model-under-training=%t -ml-inliner-ir2native-model=%S/../../../../unittests/Analysis/Inputs/ir2native_x86_64_model -S < %s | FileCheck %s 
-; RUN: opt -enable-ml-inliner=development -passes=scc-oz-module-inliner -training-log=- -ml-inliner-model-under-training=%t -ml-inliner-ir2native-model=%S/../../../../unittests/Analysis/Inputs/ir2native_x86_64_model -ml-inliner-output-spec-override=%S/Inputs/test_output_spec.json -S < %s | FileCheck %s --check-prefixes=EXTRA-OUTPUTS,CHECK
+; RUN: opt -enable-ml-inliner=development -passes=scc-oz-module-inliner -training-log=- -tfutils-text-log -ml-inliner-model-under-training=%t -ml-inliner-ir2native-model=%S/../../../../unittests/Analysis/Inputs/ir2native_x86_64_model -ml-inliner-output-spec-override=%S/Inputs/test_output_spec.json -S < %s | FileCheck %s --check-prefixes=EXTRA-OUTPUTS,CHECK
-; RUN: opt -enable-ml-inliner=development -passes=scc-oz-module-inliner -training-log=- -ml-inliner-ir2native-model=%S/../../../../unittests/Analysis/Inputs/ir2native_x86_64_model -S < %s | FileCheck %s
+; RUN: opt -enable-ml-inliner=development -passes=scc-oz-module-inliner -training-log=- -tfutils-text-log -ml-inliner-ir2native-model=%S/../../../../unittests/Analysis/Inputs/ir2native_x86_64_model -S < %s | FileCheck %s
-; RUN: opt -enable-ml-inliner=development -passes=scc-oz-module-inliner -training-log=- -ml-inliner-model-under-training=%t -S < %s | FileCheck %s --check-prefix=NOREWARD
+; RUN: opt -enable-ml-inliner=development -passes=scc-oz-module-inliner -training-log=- -tfutils-text-log -ml-inliner-model-under-training=%t -S < %s | FileCheck %s --check-prefix=NOREWARD
-; RUN: opt -enable-ml-inliner=development -passes=scc-oz-module-inliner -training-log=- -S < %s | FileCheck %s --check-prefix=NOREWARD
+; RUN: opt -enable-ml-inliner=development -passes=scc-oz-module-inliner -training-log=- -tfutils-text-log -S < %s | FileCheck %s --check-prefix=NOREWARD
-
 target datalayout = "e-m:e-i64:64-f80:128-n8:16:32:64-S128"
 target triple = "x86_64-pc-linux-gnu"
-
 declare i32 @f1(i32)
 declare i32 @f2(i32)
-
 define dso_local i32 @branches(i32) {
   %cond = icmp slt i32 %0, 3
   br i1 %cond, label %then, label %else
-
 then:
   %ret.1 = call i32 @f1(i32 %0)
   br label %last.block
-
 else:
   %ret.2 = call i32 @f2(i32 %0)
   br label %last.block
-
 last.block:
   %ret = phi i32 [%ret.1, %then], [%ret.2, %else]
   ret i32 %ret
 }
-
 define dso_local i32 @top() {
   %1 = call i32 @branches(i32 2)
   ret i32 %1
 }
-
-
 !llvm.module.flags = !{!0}
 !llvm.ident = !{!1}
-
 !0 = !{i32 1, !"wchar_size", i32 4}
 !1 = !{!"clang version 7.0.0-6 (tags/RELEASE_700/final)"}
-
 ; Check we produce a protobuf that has inlining decisions and rewards.
+; CHECK:                  key: "delta_size"
+; CHECK-NEXT:               value {
+; CHECK-NEXT:                 feature {
+; CHECK-NEXT:                   int64_list {
+; CHECK-NEXT:                     value: 0
+; CHECK-NEXT:                   }
+; CHECK-NEXT:                 }
 ; CHECK-NOT: fake_extra_output
-; EXTRA-OUTPUTS:          key: "fake_extra_output" value: {
+; EXTRA-OUTPUTS:          key: "fake_extra_output"
-; EXTRA-OUTPUTS-NEXT:       feature: { int64_list: { value: [{{[0-9]+}}] } }
+; EXTRA-OUTPUTS-NEXT:       value {
-; CHECK:          key: "inlining_decision" value: {
+; EXTRA-OUTPUTS-NEXT:         feature {
-; CHECK-NEXT:       feature: { int64_list: { value: [1] } }
+; EXTRA-OUTPUTS-NEXT:           int64_list {
-; CHECK:          key: "delta_size" value: {
+; EXTRA-OUTPUTS-NEXT:             value: {{[0-9]+}}
-; CHECK-NEXT:       feature: { int64_list: { value: [0] } }
+; CHECK:                  key: "inlining_decision"
-; CHECK-NEXT:     }
+; CHECK-NEXT:               value {
-; CHECK-NEXT:   }
+; CHECK-NEXT:                 feature {
-; NOREWARD-NOT: key: "delta_size" value: {
+; CHECK-NEXT:                   int64_list {
+; CHECK-NEXT:                     value: 1
+; NOREWARD-NOT: key: "delta_size"

unittests/Analysis/TFUtilsTest.cpp

@@ -7,6 +7,8 @@
 //===----------------------------------------------------------------------===//
 
 #include "llvm/Analysis/Utils/TFUtils.h"
+#include "tensorflow/core/example/example.pb.h"
+#include "tensorflow/core/example/feature.pb.h"
 #include "llvm/AsmParser/Parser.h"
 #include "llvm/IR/Dominators.h"
 #include "llvm/IR/Instructions.h"
@@ -143,6 +145,18 @@ TEST(TFUtilsTest, TensorSpecSizesAndTypes) {
   EXPECT_EQ(Spec1D.getElementByteSize(), sizeof(int16_t));
 }
 
+#define PROTO_CHECKER(FNAME, TYPE, INDEX, EXP)                                 \
+  do {                                                                         \
+    const auto &V = Expected.feature_lists()                                   \
+                        .feature_list()                                        \
+                        .at(FNAME)                                             \
+                        .feature(INDEX)                                        \
+                        .TYPE()                                                \
+                        .value();                                              \
+    for (auto I = 0; I < V.size(); ++I)                                        \
+      EXPECT_EQ(V.at(I), EXP[I]);                                              \
+  } while (false)
+
 TEST(TFUtilsTest, Logger) {
   std::vector<LoggedFeatureSpec> Features;
   Features.push_back(
@@ -152,42 +166,31 @@ TEST(TFUtilsTest, Logger) {
 
   auto Rewards = TensorSpec::createSpec<float>("reward", {1});
   Logger L(Features, Rewards, true);
-  float F00[]{0.0, 0.1, 0.2, 0.3, 0.4, 0.5};
+  const float F00[]{0.0, 0.1, 0.2, 0.3, 0.4, 0.5};
-  int64_t F01[]{2, 3};
+  const int64_t F01[]{2, 3};
 
   L.logTensorValue(0, F00, 6);
   L.logTensorValue(1, F01, 2);
   L.logReward<float>(3.4);
-  float F10[]{0.0, 1.0, 2.0, 3.0, 4.0, 5.0};
+  const float F10[]{0.0, 1.0, 2.0, 3.0, 4.0, 5.0};
-  int64_t F11[]{-2, -3};
+  const int64_t F11[]{-2, -3};
   L.logTensorValue(0, F10, 6);
   L.logTensorValue(1, F11, 2);
   L.logReward<float>(-3.0);
-  const auto *Expected = R"(feature_lists: {
-  feature_list: {
-    key: "the_float" value: {
-      feature: { float_list: { value: [0.000000e+00, 1.000000e-01, 2.000000e-01, 3.000000e-01, 4.000000e-01, 5.000000e-01] } }
-      feature: { float_list: { value: [0.000000e+00, 1.000000e+00, 2.000000e+00, 3.000000e+00, 4.000000e+00, 5.000000e+00] } }
-    }
-  }
-  feature_list: {
-    key: "alternate_name" value: {
-      feature: { int64_list: { value: [2, 3] } }
-      feature: { int64_list: { value: [-2, -3] } }
-    }
-  }
-  feature_list: {
-    key: "reward" value: {
-      feature: { float_list: { value: [3.400000e+00] } }
-      feature: { float_list: { value: [-3.000000e+00] } }
-    }
-  }
-}
-)";
   std::string Result;
   raw_string_ostream OS(Result);
   L.print(OS);
-  EXPECT_EQ(Result, Expected);
+
+  tensorflow::SequenceExample Expected;
+  EXPECT_TRUE(Expected.ParseFromString(Result));
+  PROTO_CHECKER("the_float", float_list, 0, F00);
+  PROTO_CHECKER("the_float", float_list, 1, F10);
+  PROTO_CHECKER("alternate_name", int64_list, 0, F01);
+  PROTO_CHECKER("alternate_name", int64_list, 1, F11);
+  float R0[]{3.4};
+  float R1[]{-3.0};
+  PROTO_CHECKER("reward", float_list, 0, R0);
+  PROTO_CHECKER("reward", float_list, 1, R1);
 }
 
 TEST(TFUtilsTest, LoggerNoReward) {
@@ -199,34 +202,25 @@ TEST(TFUtilsTest, LoggerNoReward) {
 
   auto Rewards = TensorSpec::createSpec<float>("reward", {1});
   Logger L(Features, Rewards, false);
-  float F00[]{0.0, 0.1, 0.2, 0.3, 0.4, 0.5};
+  const float F00[]{0.0, 0.1, 0.2, 0.3, 0.4, 0.5};
-  int64_t F01[]{2, 3};
+  const int64_t F01[]{2, 3};
 
   L.logTensorValue(0, F00, 6);
   L.logTensorValue(1, F01, 2);
-  float F10[]{0.0, 1.0, 2.0, 3.0, 4.0, 5.0};
+  const float F10[]{0.0, 1.0, 2.0, 3.0, 4.0, 5.0};
-  int64_t F11[]{-2, -3};
+  const int64_t F11[]{-2, -3};
   L.logTensorValue(0, F10, 6);
   L.logTensorValue(1, F11, 2);
-  const auto *Expected = R"(feature_lists: {
+
-  feature_list: {
-    key: "the_float" value: {
-      feature: { float_list: { value: [0.000000e+00, 1.000000e-01, 2.000000e-01, 3.000000e-01, 4.000000e-01, 5.000000e-01] } }
-      feature: { float_list: { value: [0.000000e+00, 1.000000e+00, 2.000000e+00, 3.000000e+00, 4.000000e+00, 5.000000e+00] } }
-    }
-  }
-  feature_list: {
-    key: "alternate_name" value: {
-      feature: { int64_list: { value: [2, 3] } }
-      feature: { int64_list: { value: [-2, -3] } }
-    }
-  }
-}
-)";
   std::string Result;
   raw_string_ostream OS(Result);
   L.print(OS);
-  EXPECT_EQ(Result, Expected);
+  tensorflow::SequenceExample Expected;
+  EXPECT_TRUE(Expected.ParseFromString(Result));
+  PROTO_CHECKER("the_float", float_list, 0, F00);
+  PROTO_CHECKER("the_float", float_list, 1, F10);
+  PROTO_CHECKER("alternate_name", int64_list, 0, F01);
+  PROTO_CHECKER("alternate_name", int64_list, 1, F11);
 }
 
 TEST(TFUtilsTest, LoggerFinalReward) {
@@ -242,32 +236,14 @@ TEST(TFUtilsTest, LoggerFinalReward) {
     L.logTensorValue(1, &I);
   }
   L.logFinalReward<float>(3.14);
-  const auto *Expected = R"(feature_lists: {
-  feature_list: {
-    key: "the_float" value: {
-      feature: { float_list: { value: [0.000000e+00] } }
-      feature: { float_list: { value: [1.000000e+00] } }
-      feature: { float_list: { value: [2.000000e+00] } }
-    }
-  }
-  feature_list: {
-    key: "the_int" value: {
-      feature: { int64_list: { value: [0] } }
-      feature: { int64_list: { value: [1] } }
-      feature: { int64_list: { value: [2] } }
-    }
-  }
-  feature_list: {
-    key: "reward" value: {
-      feature: { float_list: { value: [0.000000e+00] } }
-      feature: { float_list: { value: [0.000000e+00] } }
-      feature: { float_list: { value: [3.140000e+00] } }
-    }
-  }
-}
-)";
   std::string Result;
   raw_string_ostream OS(Result);
   L.print(OS);
-  EXPECT_EQ(Result, Expected);
+  const float Zero[]{0.0};
+  const float R[]{3.14};
+  tensorflow::SequenceExample Expected;
+  EXPECT_TRUE(Expected.ParseFromString(Result));
+  PROTO_CHECKER("reward", float_list, 0, Zero);
+  PROTO_CHECKER("reward", float_list, 1, Zero);
+  PROTO_CHECKER("reward", float_list, 2, R);
 }