Merge branch 'panoptic-segmentation' into panoptic-segmentation

research/seq_flow_lite/tf_ops/projection_util.h

@@ -21,23 +21,26 @@ limitations under the License.
 
 #include "libutf/utf.h"
 
-constexpr int kFirstCapOffset = 3;
-constexpr int kAllCapsOffset = 4;
-constexpr int kWordNoveltyOffset = 1;
-constexpr int kDocSizeOffset = 2;
-
-const char kMurmurHash[] = "murmur";
-const char kXfixHash8[] = "xfixhash8";
-const char kXfixHash16[] = "xfixhash16";
-const char kXfixHash32[] = "xfixhash32";
-const char kUnicodeHash8[] = "unicodehash8";
-const char kUnicodeHash16[] = "unicodehash16";
-
+inline constexpr int kFirstCapOffset = 3;
+inline constexpr int kAllCapsOffset = 4;
+inline constexpr int kWordNoveltyOffset = 1;
+inline constexpr int kDocSizeOffset = 2;
+
+inline constexpr char kMurmurHash[] = "murmur";
+inline constexpr char kXfixHash8[] = "xfixhash8";
+inline constexpr char kXfixHash16[] = "xfixhash16";
+inline constexpr char kXfixHash32[] = "xfixhash32";
+inline constexpr char kUnicodeHash8[] = "unicodehash8";
+inline constexpr char kUnicodeHash16[] = "unicodehash16";
+
+// A base class that specifies the interface for a hash engine used by the
+// projection operator.
 class HashEngine {
  public:
-  virtual void GetHashCodes(const std::string& word,
-                            std::vector<uint64_t>* hash_codes,
-                            int feature_size) = 0;
+  // Takes a string token `word` and a `feature_size` (measured in bits) and
+  // returns hash codes that represent the token.
+  virtual std::vector<uint64_t> GetHashCodes(const std::string& word,
+                                             int feature_size) = 0;
   virtual ~HashEngine() {}
 };
 
@@ -50,13 +53,12 @@ class Hasher {
                               const std::string& hash_type = kMurmurHash);
   static bool SupportedHashType(const std::string& hash_type);
   bool GetHashCodes(const std::string& word,
-                    std::vector<uint64_t>* hash_codes) {
+                    std::vector<uint64_t>& hash_codes) {
     if (!hash_engine_) return false;
     if (word.empty()) {
-      *hash_codes = null_hash_codes_;
+      hash_codes = null_hash_codes_;
     } else {
-      hash_codes->clear();
-      hash_engine_->GetHashCodes(word, hash_codes, feature_size_);
+      hash_codes = hash_engine_->GetHashCodes(word, feature_size_);
     }
     return true;
   }
@@ -64,8 +66,13 @@ class Hasher {
  private:
   explicit Hasher(int feature_size, HashEngine* hash_engine);
   const std::string empty_string_ = "<null>";
+  // Size of the projection feature which represents the number of bits of
+  // hash codes that will be generated by this class.
   const int feature_size_;
+  // The hash engine used by this class.
   std::unique_ptr<HashEngine> hash_engine_;
+  // Hash codes for empty text is precaculated and stored below to speed
+  // up projection.
   std::vector<uint64_t> null_hash_codes_;
 };
 
@@ -90,7 +97,8 @@ class ProjectionUnicodeHandler {
   }
 
   // Performs language independent lower case and returns a string with
-  // supported unicode segments.
+  // supported unicode segments and two additional flags first_cap and all_caps
+  // which when true indicate the text is Firstcap or ALLCAPS.
   std::string LowerCaseUTF8WithSupportedUnicodes(
       const std::pair<const char*, size_t>& source, bool* first_cap = nullptr,
       bool* all_caps = nullptr) const;
@@ -126,14 +134,19 @@ class ProjectionUnicodeHandler {
                                                               int max_tokens);
 
  private:
-  // Parses and extracts supported unicode segments from a utf8 string.
+  // Parses and extracts supported or allowed unicode segments, also referred
+  // to as vocabulary, from a utf8 string.
   void InitializeVocabulary(const std::string& vocabulary);
+  // A variable that maps a valid Unicode rune to its index in valid character
+  // vocabulary.
   std::unordered_map<Rune, int> valid_chars_;
+  // Controls whether to exclude non-alphabetic, non-space characters from the
+  // output text.
   bool exclude_nonalphaspace_unicodes_;
 };
 
-static constexpr size_t kEntireString = SIZE_MAX;
-static constexpr size_t kAllTokens = SIZE_MAX;
+inline constexpr size_t kEntireString = SIZE_MAX;
+inline constexpr size_t kAllTokens = SIZE_MAX;
 
 std::vector<std::string> SplitBySpace(const char* input_ptr, size_t len,
                                       size_t max_input, size_t max_tokens);

research/seq_flow_lite/tf_ops/repo.bzl

@@ -198,7 +198,7 @@ cc_library(
     ),
     includes = ["tensorflow_includes"],
     deps = [
-        "@eigen_archive//:eigen",
+        "@eigen_archive//:eigen3",
         "@protobuf_archive//:includes",
         "@zlib_includes//:includes",
         "@snappy_includes//:includes",

research/seq_flow_lite/tf_ops/sequence_string_projection.cc

@@ -12,16 +12,16 @@ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
-#include "tf_ops/projection_normalizer_util.h"  // seq_flow_lite
-#include "tf_ops/projection_tokenizer_util.h"  // seq_flow_lite
-#include "tf_ops/projection_util.h"  // seq_flow_lite
-#include "tf_ops/text_distorter.h"  // seq_flow_lite
 #include "absl/container/flat_hash_map.h"
 #include "tensorflow/core/framework/op.h"
 #include "tensorflow/core/framework/op_kernel.h"
 #include "tensorflow/core/framework/shape_inference.h"
 #include "tensorflow/core/framework/tensor.h"
 #include "tensorflow/core/framework/tensor_shape.h"
+#include "tf_ops/projection_normalizer_util.h"  // seq_flow_lite
+#include "tf_ops/projection_tokenizer_util.h"  // seq_flow_lite
+#include "tf_ops/projection_util.h"  // seq_flow_lite
+#include "tf_ops/text_distorter.h"  // seq_flow_lite
 
 using ::tensorflow::int32;
 using ::tensorflow::int64;
@@ -51,10 +51,11 @@ float* AllocateTensor(OpKernelContext* ctx, const std::string& tensor_name,
   return &tensor->flat<float>()(0);
 }
 
+// OpKernel for the sequence string projection op.
 class SequenceStringProjectionOp : public OpKernel {
  public:
   explicit SequenceStringProjectionOp(OpKernelConstruction* context)
-      : OpKernel(context) {
+      : OpKernel(context), philox_(171), generator_(&philox_) {
     OP_REQUIRES_OK(context, context->GetAttr("feature_size", &feature_size_));
     std::string hashtype;
     OP_REQUIRES_OK(context, context->GetAttr("hashtype", &hashtype));
@@ -159,7 +160,10 @@ class SequenceStringProjectionOp : public OpKernel {
       }
       const int64 seq_len =
           static_cast<int64>(bos_tag_ + words.size() + eos_tag_);
-      CHECK_GT(seq_len, 0);
+      CHECK_GT(seq_len, 0)
+          << "Projection models expect input text to have at-least one valid "
+             "token. If empty text is a valid input for your model, please set "
+             "add_bos_tag to true.";
       max_seq_len = std::max(max_seq_len, seq_len);
       words_batches.emplace_back(std::move(words));
     }
@@ -208,7 +212,7 @@ class SequenceStringProjectionOp : public OpKernel {
           CHECK_EQ(eos_tag_, 1);
           word = kEndTokenTSP;
         }
-        hasher_->GetHashCodes(word, &hash_codes);
+        hasher_->GetHashCodes(word, hash_codes);
         for (int hindex = 0, k = 0; hindex < hash_codes.size(); hindex++) {
           auto hash = hash_codes[hindex];
           for (int kmax = std::min(k + increment, feature_size_); k < kmax;) {
@@ -229,7 +233,7 @@ class SequenceStringProjectionOp : public OpKernel {
               doc_size_feature;
         }
         if (add_first_cap_feature_ > 0.0f) {
-          if (text_distorter_->BernouilleSample(add_first_cap_feature_)) {
+          if (generator_.RandFloat() <= add_first_cap_feature_) {
             projection[offset0 + feature_size_ - kFirstCapOffset] =
                 first_cap ? 1.0 : -1.0;
           } else {
@@ -237,7 +241,7 @@ class SequenceStringProjectionOp : public OpKernel {
           }
         }
         if (add_all_caps_feature_ > 0.0f) {
-          if (text_distorter_->BernouilleSample(add_all_caps_feature_)) {
+          if (generator_.RandFloat() <= add_all_caps_feature_) {
             projection[offset0 + feature_size_ - kAllCapsOffset] =
                 all_caps ? 1.0 : -1.0;
           } else {
@@ -252,21 +256,49 @@ class SequenceStringProjectionOp : public OpKernel {
   }
 
  private:
+  // Objects used for random number generator.
+  tensorflow::random::PhiloxRandom philox_;
+  tensorflow::random::SimplePhilox generator_;
+
+  // Dimensionality of the ternary vector for each token in the text.
   int32 feature_size_;
+  // An object used to hash tokens in the text.
   std::unique_ptr<Hasher> hasher_;
+  // An object used for distorting text before projection.
   std::unique_ptr<TextDistorter> text_distorter_;
+  // An object used for manipulating unicode in the text. It performs tasks such
+  // as retaining only whitelisted unicodes in the text tokens and lowercasing
+  // them.
   std::unique_ptr<ProjectionUnicodeHandler> unicode_handler_;
+  // An object used for normalizing tokens in the text. This performs tasks
+  // such as identifying repeated characters and replace them with a single
+  // instance.
   std::unique_ptr<ProjectionNormalizer> projection_normalizer_;
+  // Character whitelist used by the projection operator.
   std::string vocabulary_;
+  // Size of the character whitelist.
   int vocabulary_size_;
+  // Maximum number of splits allowed in the text. The number of tokens in text
+  // post segmentation will be utmost max_splits_ + 1.
   int32 max_splits_;
+  // A flag that indicates how to segment text. When true text is segmented by
+  // space. Otherwise it is segmented on unicode boundaries.
   bool split_on_space_;
+  // When true include an end of sentence token in the projection.
   int eos_tag_;
+  // When true include a begin of sentence token in the projection.
   int bos_tag_;
+  // Number of bits used to capture word novelty. See tensorflow op
+  // documentation below for details.
   int word_novelty_bits_;
+  // Number of levels used to capture document size. See tensorflow op
+  // documentation below for details.
   int doc_size_levels_;
+  // Distance between levels used for word novelty.
   float word_novelty_offset_;
+  // Adds boolean feature to indicate first_cap text with the below probability.
   float add_first_cap_feature_;
+  // Adds boolean feature to indicate all_cap text with the below probability.
   float add_all_caps_feature_;
 };
 

research/seq_flow_lite/tf_ops/sequence_string_projection_op_v2.cc

@@ -40,6 +40,8 @@ constexpr char kEndTokenTSP[] = "<EOS>";
 constexpr float kMappingTable[4] = {0, 1, -1, 0};
 constexpr int kIncrement = 32;
 
+// Version 2 OpKernel for the sequence string projection op.
+// Template T can be int32 or int64.
 template <typename T>
 class SequenceStringProjectionOpV2 : public OpKernel {
  public:
@@ -136,7 +138,7 @@ class SequenceStringProjectionOpV2 : public OpKernel {
         } else {
           word = kEndTokenTSP;
         }
-        hasher_->GetHashCodes(word, &hash_codes);
+        hasher_->GetHashCodes(word, hash_codes);
         for (int hindex = 0, k = 0; hindex < hash_codes.size(); hindex++) {
           auto hash = hash_codes[hindex];
           for (int kmax = std::min(k + kIncrement, feature_size_); k < kmax;) {
@@ -153,13 +155,25 @@ class SequenceStringProjectionOpV2 : public OpKernel {
   }
 
  private:
+  // Dimensionality of the ternary vector for each token in the text.
   int32 feature_size_;
+  // An object used to hash tokens in the text.
   std::unique_ptr<Hasher> hasher_;
+  // An object used for distorting text before projection.
   std::unique_ptr<TextDistorter> text_distorter_;
+  // An object used for manipulating unicode in the text. It performs tasks such
+  // as retaining only whitelisted unicodes in the text tokens and lowercasing
+  // them.
   std::unique_ptr<ProjectionUnicodeHandler> unicode_handler_;
+  // An object used for normalizing tokens in the text. This performs tasks
+  // such as identifying repeated characters and replace them with a single
+  // instance.
   std::unique_ptr<ProjectionNormalizer> projection_normalizer_;
+  // Character whitelist used by the projection operator.
   std::string vocabulary_;
+  // When true include an end of sentence token in the projection.
   int eos_tag_;
+  // When true include a begin of sentence token in the projection.
   int bos_tag_;
 };
 

research/seq_flow_lite/tf_ops/text_distorter.h

@@ -32,7 +32,6 @@ class TextDistorter {
     assert(distortion_probability_ <= 1.0);
   }
   std::string DistortText(icu::UnicodeString* uword);
-  bool BernouilleSample(float p) { return (generator_.RandFloat() <= p); }
 
  private:
   tensorflow::random::PhiloxRandom philox_;

research/seq_flow_lite/tf_ops/tf_custom_ops.cc

@@ -20,30 +20,6 @@ limitations under the License.
 
 using ::tensorflow::int32;
 
-class PoolingOp : public tensorflow::OpKernel {
- public:
-  explicit PoolingOp(tensorflow::OpKernelConstruction* context)
-      : tensorflow::OpKernel(context) {}
-
-  void Compute(tensorflow::OpKernelContext* ctx) override {}
-};
-
-REGISTER_KERNEL_BUILDER(Name("PoolingOp").Device(::tensorflow::DEVICE_CPU),
-                        PoolingOp);
-
-REGISTER_OP("PoolingOp")
-    .Input("multiplier: float32")
-    .Input("constant: float32")
-    .Input("forward: float32")
-    .Output("state: float32")
-    .SetShapeFn([](::tensorflow::shape_inference::InferenceContext* c) {
-      c->set_output(0, c->input(0));
-      return tensorflow::Status::OK();
-    })
-    .Doc(R"doc(
-Dummy pooling op.
-)doc");
-
 class ExpectedValueOp : public tensorflow::OpKernel {
  public:
   explicit ExpectedValueOp(tensorflow::OpKernelConstruction* context)
@@ -93,33 +69,3 @@ REGISTER_OP("LayerNorm")
     .Doc(R"doc(
 Dummy layer norm op.
 )doc");
-
-class UniformCausalAttnOp : public tensorflow::OpKernel {
- public:
-  explicit UniformCausalAttnOp(tensorflow::OpKernelConstruction* context)
-      : tensorflow::OpKernel(context) {}
-
-  void Compute(tensorflow::OpKernelContext* ctx) override {}
-};
-
-REGISTER_KERNEL_BUILDER(
-    Name("UniformCausalAttn").Device(::tensorflow::DEVICE_CPU),
-    UniformCausalAttnOp);
-
-REGISTER_OP("UniformCausalAttn")
-    .Input("input: float32")
-    .Input("time_step: int32")
-    .Input("selected_beams: int32")
-    .Attr("feature_size: int")
-    .Attr("beam_size: int")
-    .Output("output: float32")
-    .SetShapeFn([](::tensorflow::shape_inference::InferenceContext* c) {
-      auto batch_size = c->Dim(c->input(0), 0);
-      int32 feature_size;
-      TF_RETURN_IF_ERROR(c->GetAttr("feature_size", &feature_size));
-      c->set_output(0, c->MakeShape({batch_size, 1, feature_size}));
-      return tensorflow::Status::OK();
-    })
-    .Doc(R"doc(
-Dummy uniform causal attn op.
-)doc");

research/seq_flow_lite/tflite_ops/BUILD

 # TFLite ops for sequence string projection.
+load("@org_tensorflow//tensorflow:tensorflow.bzl", "pybind_extension")
 load("@org_tensorflow//tensorflow/lite:build_def.bzl", "tflite_copts")
 
 licenses(["notice"])
@@ -100,3 +101,16 @@ cc_test(
         "@flatbuffers",
     ],
 )
+
+pybind_extension(
+    name = "registerer",
+    srcs = ["registerer.cc"],
+    module_name = "registerer",
+    deps = [
+        ":expected_value",
+        ":layer_norm",
+        ":sequence_string_projection",
+        "@org_tensorflow//tensorflow/lite:framework",
+        "@pybind11",
+    ],
+)

research/seq_flow_lite/tflite_ops/expected_value.cc

@@ -18,7 +18,7 @@ limitations under the License.
 
 #include "tflite_ops/quantization_util.h"  // seq_flow_lite
 
-namespace tflite {
+namespace seq_flow_lite {
 namespace ops {
 namespace custom {
 
@@ -156,4 +156,4 @@ TfLiteRegistration* Register_EXPECTED_VALUE() {
 
 }  // namespace custom
 }  // namespace ops
-}  // namespace tflite
+}  // namespace seq_flow_lite

research/seq_flow_lite/tflite_ops/expected_value.h

@@ -17,7 +17,7 @@ limitations under the License.
 
 #include "tensorflow/lite/kernels/register.h"
 
-namespace tflite {
+namespace seq_flow_lite {
 namespace ops {
 namespace custom {
 
@@ -25,6 +25,6 @@ TfLiteRegistration* Register_EXPECTED_VALUE();
 
 }  // namespace custom
 }  // namespace ops
-}  // namespace tflite
+}  // namespace seq_flow_lite
 
 #endif  // TENSORFLOW_MODELS_SEQUENCE_PROJECTION_TFLITE_OPS_EXPECTED_VALUE_H_

research/seq_flow_lite/tflite_ops/layer_norm.cc

@@ -17,10 +17,10 @@ limitations under the License.
 #include <unordered_set>
 #include <vector>
 
-#include "tflite_ops/quantization_util.h"  // seq_flow_lite
 #include "tensorflow/lite/kernels/kernel_util.h"
+#include "tflite_ops/quantization_util.h"  // seq_flow_lite
 
-namespace tflite {
+namespace seq_flow_lite {
 namespace ops {
 namespace custom {
 
@@ -213,6 +213,102 @@ TfLiteStatus FlexibleLayerNorm(const TfLiteTensor* input, const float scale,
   return kTfLiteOk;
 }
 
+/*
+ * Layer normalization is optimized as follows in integer arithmetic
+ *
+ * Algorithm
+ * *********
+ * Subscript i \in {1, ..., N}, Inputs q_i, Outputs oq_i.
+ *
+ * x_i = (q_i - input_zero_point) * input_scale
+ * mean = sum_i x_i / N
+ * var = sum_i (x_i * x_i / N) - mean * mean
+ * std = sqrt(var + tolerance)
+ * xni = (xi - mean) / std
+ * yi = xni * scale + offset
+ * o_i = round(y_i / output_scale + output_zero_point)
+ * oq_i = clamp(o_i, 0, 255)
+ *
+ * Optimizations
+ * *************
+ * Applying linear expansion
+ * x_i = q_i * input_scale - input_zero_point * input_scale
+ * or x_i = m * qi + c
+ * mean = m * mean_q + c
+ * Variance is not affected by a constant shift to input
+ * var = m^2 * var_q
+ * std = m * sqrt(var_q + tolerance)
+ * Expanding xi, mean, std in equation for xni
+ * xni = (m * qi + c - m * mean_q - c) / m * sqrt(var_q + tolerance)
+ * Simplifying
+ * xni = (qi - mean_q) / sqrt(var_q + tolerance)
+ * Setting inv_std_qi = 1 / sqrt(var_q + tolerance)
+ * xni = qi * inv_std_qi - mean_q * inv_std_qi
+ * yi = qi * inv_std_qi * scale - mean_q * inv_std_qi * scale + offset
+ * o_i = round(qi * inv_std_qi * scale / output_scale
+ *             - mean_q * inv_std_qi * scale / output_scale
+ *             + offset / output_scale
+ *             + output_zero_point)
+ * Setting
+ * static_bias = offset / output_scale + output_zero_point
+ * static_scale = scale / output_scale
+ * o_i = round(qi * inv_std_qi * static_scale
+ *             - mean_q * inv_std_qi * static_scale
+ *             + static_bias)
+ * Setting
+ * dynamic_scale = inv_std_qi * static_scale
+ * dynamic_bias = static_bias - mean_q * dynamic_scale
+ * o_i = round(qi * dynamic_scale + dynamic_bias)
+ * oq_i = clamp(round(qi * dynamic_scale + dynamic_bias), 0, 255)
+ *
+ * This results in the below optimized implementation. The strategy is to first
+ * compute first and second order summary statistics for qi in a loop,
+ * then compute mean_q, var_q and then dynamic_scale/dynamic_bias. This
+ * allows one to compute oqi quickly in a tight loop.
+ * */
+TfLiteStatus IntegerLayerNorm(const TfLiteTensor* input, const float scale,
+                              const float offset, TfLiteTensor* output) {
+  const int input_rank = input->dims->size;
+  const int num_features = input->dims->data[input_rank - 1];
+  const int time_steps =
+      static_cast<int>(GetNumberOfSteps(input) / num_features);
+
+  const float out_inverse_scale = 1.0f / output->params.scale;
+  const float static_scale = scale * out_inverse_scale;
+  const float static_bias = static_cast<float>(output->params.zero_point) +
+                            offset * out_inverse_scale;
+  const float inverse_num_features = 1.0f / num_features;
+  const uint8_t* const in_ptr = input->data.uint8;
+  uint8_t* out_ptr = output->data.uint8;
+  for (int i = 0; i < time_steps; ++i) {
+    int32_t i32_sum_q = 0;
+    int32_t i32_sum_qq = 0;
+    const int32_t index = i * num_features;
+    for (int j = index; j < index + num_features; ++j) {
+      const int32_t q_i = static_cast<int32_t>(in_ptr[j]);
+      // Compute first and second order statistics for qi.
+      i32_sum_q += q_i;
+      i32_sum_qq += q_i * q_i;
+    }
+    const float second_moment_qq = i32_sum_qq * inverse_num_features;
+    const float mean_q = i32_sum_q * inverse_num_features;
+    const float var_q = second_moment_qq - mean_q * mean_q;
+    const float inv_std_q = 1.0f / sqrt(var_q + 1e-6);
+    const float dynamic_scale = inv_std_q * static_scale;
+    const float dynamic_bias = static_bias - mean_q * dynamic_scale;
+    for (int j = index; j < index + num_features; ++j) {
+      const int32_t invalue = static_cast<int32_t>(in_ptr[j]);
+      const float value = invalue * dynamic_scale + dynamic_bias;
+      // Use an offseted cast to perform float round.
+      const int32_t i32value =
+          static_cast<int32_t>(value + ((value >= 0.0) ? 0.5f : -0.5f));
+      // Clamp the result.
+      out_ptr[j] = static_cast<uint8_t>(std::max(std::min(255, i32value), 0));
+    }
+  }
+  return kTfLiteOk;
+}
+
 TfLiteStatus DefaultLayerNormFloat(const TfLiteTensor* input, const float scale,
                                    const float offset, TfLiteTensor* output) {
   const int input_rank = input->dims->size;
@@ -298,7 +394,7 @@ TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
   if (num_axis == 1 && (axis->data.i32[0] == -1 ||
                         axis->data.i32[0] == (input->dims->size - 1))) {
     if (input->type == kTfLiteUInt8) {
-      return DefaultLayerNorm(input, scale, offset, output);
+      return IntegerLayerNorm(input, scale, offset, output);
     } else if (input->type == kTfLiteFloat32) {
       return DefaultLayerNormFloat(input, scale, offset, output);
     } else {
@@ -328,4 +424,4 @@ TfLiteRegistration* Register_LAYER_NORM() {
 
 }  // namespace custom
 }  // namespace ops
-}  // namespace tflite
+}  // namespace seq_flow_lite

research/seq_flow_lite/tflite_ops/layer_norm.h

@@ -17,7 +17,7 @@ limitations under the License.
 
 #include "tensorflow/lite/kernels/register.h"
 
-namespace tflite {
+namespace seq_flow_lite {
 namespace ops {
 namespace custom {
 
@@ -25,6 +25,6 @@ TfLiteRegistration* Register_LAYER_NORM();
 
 }  // namespace custom
 }  // namespace ops
-}  // namespace tflite
+}  // namespace seq_flow_lite
 
 #endif  // LEARNING_EXPANDER_POD_DEEP_POD_TFLITE_HANDLERS_LAYER_NORM_H_

research/seq_flow_lite/tflite_ops/layer_norm_test.cc

@@ -20,40 +20,35 @@ limitations under the License.
 #include "flatbuffers/flexbuffers.h"  // flatbuffer
 #include "tensorflow/lite/kernels/test_util.h"
 
-namespace tflite {
+namespace seq_flow_lite {
 namespace ops {
 namespace custom {
 namespace {
 
-class LayerNormModel : public SingleOpModel {
+using ::testing::ElementsAreArray;
+using ::tflite::ArrayFloatNear;
+using ::tflite::Dequantize;
+using ::tflite::TensorType_INT32;
+using ::tflite::TensorType_UINT8;
+
+class LayerNormModel : public ::tflite::SingleOpModel {
  public:
-  explicit LayerNormModel(const TensorData& input, float output_min,
+  explicit LayerNormModel(std::initializer_list<int> input_shape,
+                          float input_min, float input_max, float output_min,
                           float output_max, float scale, float offset,
-                          std::initializer_list<int> axis_shape,
-                          std::initializer_list<int> axis)
+                          std::initializer_list<int> axes)
       : scale_value_(scale), offset_value_(offset) {
-    input_ = AddInput(input);
+    const int num_axes = axes.size();
+    input_ = AddInput({TensorType_UINT8, input_shape, input_min, input_max});
     scale_ = AddInput(
         {TensorType_UINT8, {1}, std::min(scale, 0.0f), std::max(scale, 0.0f)});
     offset_ = AddInput({TensorType_UINT8,
                         {1},
                         std::min(offset, 0.0f),
                         std::max(offset, 0.0f)});
-    axis_ = AddConstInput(TensorType_INT32, axis, axis_shape);
+    axis_ = AddConstInput(TensorType_INT32, axes, {num_axes});
     output_ = AddOutput({TensorType_UINT8, {}, output_min, output_max});
-    flexbuffers::Builder fbb;
-    fbb.Map([&] {
-      {
-        size_t start = fbb.StartVector("axes");
-        for (const int& aval : axis) {
-          fbb.Int(aval);
-        }
-        fbb.EndVector(start, /*typed=*/true, /*fixed=*/false);
-      }
-    });
-    fbb.Finish();
-
-    SetCustomOp("LayerNorm", fbb.GetBuffer(), Register_LAYER_NORM);
+    SetCustomOp("LayerNorm", {}, Register_LAYER_NORM);
     BuildInterpreter({GetShape(input_)});
   }
 
@@ -88,8 +83,9 @@ TEST(LayerNormModelTest, RegularInput) {
   const std::vector<float> expected_output = {0.0, -1.6,  0.53, 1.07,
                                               0.0, -1.13, 1.59, -0.45};
 
-  LayerNormModel m({TensorType_UINT8, {1, 2, 4}, -10, 10}, -10, 10, 1.0, 0.0,
-                   {1}, {2});
+  LayerNormModel m(/*input_shape=*/{1, 2, 4}, /*input_min=*/-10,
+                   /*input_max=*/10, /*output_min=*/-10, /*output_max=*/10,
+                   /*scale=*/1.0, /*offset=*/0.0, /*axes=*/{2});
   m.SetInput(input);
   m.Invoke();
   EXPECT_THAT(
@@ -106,8 +102,9 @@ TEST(LayerNormModelTest, NegativeScale) {
   // Standard deviation values are 3.74, 4.41
   const std::vector<float> expected_output = {0.0, 1.6,  -0.53, -1.07,
                                               0.0, 1.13, -1.59, 0.45};
-  LayerNormModel m({TensorType_UINT8, {1, 2, 4}, -10, 10}, -10, 10, -1.0, 0.0,
-                   {1}, {2});
+  LayerNormModel m(/*input_shape=*/{1, 2, 4}, /*input_min=*/-10,
+                   /*input_max=*/10, /*output_min=*/-10, /*output_max=*/10,
+                   /*scale=*/-1.0, /*offset=*/0.0, /*axes=*/{2});
   m.SetInput(input);
   m.Invoke();
   EXPECT_THAT(
@@ -124,8 +121,9 @@ TEST(LayerNormModelTest, NegativeOffset) {
   // Standard deviation values are 3.74, 4.41
   const std::vector<float> expected_output = {-1.0, -2.6,  -0.53, 0.07,
                                               -1.0, -2.13, 0.59,  -1.45};
-  LayerNormModel m({TensorType_UINT8, {1, 2, 4}, -10, 10}, -10, 10, 1.0, -1.0,
-                   {1}, {2});
+  LayerNormModel m(/*input_shape=*/{1, 2, 4}, /*input_min=*/-10,
+                   /*input_max=*/10, /*output_min=*/-10, /*output_max=*/10,
+                   /*scale=*/1.0, /*offset=*/-1.0, /*axes=*/{2});
   m.SetInput(input);
   m.Invoke();
   EXPECT_THAT(
@@ -142,8 +140,9 @@ TEST(LayerNormModelTest, NegativeScaleAndOffset) {
   // Standard deviation values are 3.74, 4.41
   const std::vector<float> expected_output = {-1.0, 0.6,  -1.53, -2.07,
                                               -1.0, 0.13, -2.59, -0.55};
-  LayerNormModel m({TensorType_UINT8, {1, 2, 4}, -10, 10}, -10, 10, -1.0, -1.0,
-                   {1}, {2});
+  LayerNormModel m(/*input_shape=*/{1, 2, 4}, /*input_min=*/-10,
+                   /*input_max=*/10, /*output_min=*/-10, /*output_max=*/10,
+                   /*scale=*/-1.0, /*offset=*/-1.0, /*axes=*/{2});
   m.SetInput(input);
   m.Invoke();
   EXPECT_THAT(
@@ -160,8 +159,9 @@ TEST(LayerNormModelTest, MultipleAxis) {
       1.12,  -2.08, 0.48,  -0.16, -0.95, -1.46, -0.95, 0.06,
       -0.69, -0.23, -1.60, -1.15, -0.80, -0.16, 0.48,  1.12};
 
-  LayerNormModel m({TensorType_UINT8, {1, 2, 3, 4}, -3, 3}, -3, 3, 1.0, 0.0,
-                   {2}, {1, 3});
+  LayerNormModel m(/*input_shape=*/{1, 2, 3, 4}, /*input_min=*/-3,
+                   /*input_max=*/3, /*output_min=*/-3, /*output_max=*/3,
+                   /*scale=*/1.0, /*offset=*/0.0, /*axes=*/{1, 3});
   m.SetInput(input);
   m.Invoke();
   EXPECT_THAT(
@@ -178,8 +178,9 @@ TEST(LayerNormModelTest, MultipleNegativeAxis) {
       1.12,  -2.08, 0.48,  -0.16, -0.95, -1.46, -0.95, 0.06,
       -0.69, -0.23, -1.60, -1.15, -0.80, -0.16, 0.48,  1.12};
 
-  LayerNormModel m({TensorType_UINT8, {1, 2, 3, 4}, -3, 3}, -3, 3, 1.0, 0.0,
-                   {2}, {-3, -1});
+  LayerNormModel m(/*input_shape=*/{1, 2, 3, 4}, /*input_min=*/-3,
+                   /*input_max=*/3, /*output_min=*/-3, /*output_max=*/3,
+                   /*scale=*/1.0, /*offset=*/0.0, /*axes=*/{-3, -1});
   m.SetInput(input);
   m.Invoke();
   EXPECT_THAT(
@@ -199,8 +200,9 @@ TEST(LayerNormModelTest, MultipleAxisWithLargeDepth) {
       2.05,  2.05,  -0.67, -0.28, 1.27,  1.27,  -1.06, -1.06, -0.28,
       0.,    -0.85, -0.42, 0.,    0.42,  -0.85, -0.42, 0.,    0.42};
 
-  LayerNormModel m({TensorType_UINT8, {1, 2, 2, 9}, -1.0, 1.0}, -3.0, 3.0, 1.0,
-                   0.0, {2}, {1, 3});
+  LayerNormModel m(/*input_shape=*/{1, 2, 2, 9}, /*input_min=*/-1.0,
+                   /*input_max=*/1.0, /*output_min=*/-3.0, /*output_max=*/3.0,
+                   /*scale=*/1.0, /*offset=*/0.0, /*axes=*/{1, 3});
   m.SetInput(input);
   m.Invoke();
   EXPECT_THAT(
@@ -211,4 +213,4 @@ TEST(LayerNormModelTest, MultipleAxisWithLargeDepth) {
 }  // namespace
 }  // namespace custom
 }  // namespace ops
-}  // namespace tflite
+}  // namespace seq_flow_lite

research/seq_flow_lite/tflite_ops/quantization_util.h

@@ -20,7 +20,7 @@ limitations under the License.
 
 #include "tensorflow/lite/context.h"
 
-namespace tflite {
+namespace seq_flow_lite {
 
 // Returns the original (dequantized) value of 8bit value.
 inline float PodDequantizeValue(const TfLiteTensor& tensor, uint8_t value) {
@@ -48,6 +48,6 @@ inline uint8_t PodQuantize(float value, int32_t zero_point,
   return static_cast<uint8_t>(std::max(std::min(255, integer_value), 0));
 }
 
-}  // namespace tflite
+}  // namespace seq_flow_lite
 
 #endif  // TENSORFLOW_MODELS_SEQUENCE_PROJECTION_TFLITE_OPS_QUANTIZATION_UTIL_H_

research/seq_flow_lite/tflite_ops/registerer.cc

+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#include "pybind11/pybind11.h"
+#include "pybind11/pytypes.h"
+#include "tensorflow/lite/mutable_op_resolver.h"
+#include "tflite_ops/expected_value.h"  // seq_flow_lite
+#include "tflite_ops/layer_norm.h"  // seq_flow_lite
+#include "tflite_ops/sequence_string_projection.h"  // seq_flow_lite
+
+PYBIND11_MODULE(registerer, m) {
+  m.doc() =
+      "Module that provides a registerer from the seq flow lite custom ops";
+  m.def(
+      "RegisterCustomOps",
+      [](uintptr_t ptr) {
+        ::tflite::MutableOpResolver* resolver =
+            reinterpret_cast<::tflite::MutableOpResolver*>(ptr);
+        resolver->AddCustom(
+            "ExpectedValueOp",
+            ::seq_flow_lite::ops::custom::Register_EXPECTED_VALUE());
+        resolver->AddCustom(
+            "LayerNorm", ::seq_flow_lite::ops::custom::Register_LAYER_NORM());
+        resolver->AddCustom("SEQUENCE_STRING_PROJECTION",
+                            ::seq_flow_lite::ops::custom::
+                                Register_SEQUENCE_STRING_PROJECTION());
+        resolver->AddCustom("SequenceStringProjection",
+                            ::seq_flow_lite::ops::custom::
+                                Register_SEQUENCE_STRING_PROJECTION());
+        resolver->AddCustom("SEQUENCE_STRING_PROJECTION_V2",
+                            ::seq_flow_lite::ops::custom::
+                                Register_SEQUENCE_STRING_PROJECTION());
+        resolver->AddCustom("SequenceStringProjectionV2",
+                            ::seq_flow_lite::ops::custom::
+                                Register_SEQUENCE_STRING_PROJECTION_V2());
+      },
+      "Register custom ops used by seq flow lite layers");
+}

research/seq_flow_lite/tflite_ops/sequence_string_projection.cc

@@ -31,7 +31,7 @@ limitations under the License.
 #include "tf_ops/projection_util.h"  // seq_flow_lite
 #include "tflite_ops/quantization_util.h"  // seq_flow_lite
 
-namespace tflite {
+namespace seq_flow_lite {
 namespace ops {
 namespace custom {
 
@@ -163,7 +163,7 @@ class ProjectionParams {
     DocSizeFeature(&doc_size_feature, num_tokens);
     *data = PodQuantize(doc_size_feature, 127.0f, 127);
   }
-  void Hash(const std::string& word, std::vector<uint64_t>* hash_codes) {
+  void Hash(const std::string& word, std::vector<uint64_t>& hash_codes) {
     hasher_->GetHashCodes(word, hash_codes);
   }
   // Lower cases the input text and eliminates all unsupported
@@ -269,6 +269,8 @@ class ProjectionParamsV2 : public ProjectionParams {
                            num_tokens, dims->data[1]);
       return kTfLiteError;
     }
+    tokens_.clear();
+    tokens_.reserve(num_tokens);
     for (int i = 0; i < num_tokens; ++i) {
       const tflite::StringRef strref = tflite::GetString(input_t, i);
       tokens_.push_back(std::pair<const char*, size_t>(strref.str, strref.len));
@@ -412,7 +414,7 @@ void TypedEval(const T* mapping_table, ProjectionParams* params, T* data) {
     } else {
       word = kEndToken;
     }
-    params->Hash(word, &hash_codes);
+    params->Hash(word, hash_codes);
     for (int hindex = 0, k = 0; hindex < hash_codes.size(); hindex++) {
       auto hash = hash_codes[hindex];
       for (int kmax = std::min(k + kIncrement, params->FeatureSize());
@@ -505,4 +507,4 @@ TfLiteRegistration* Register_SEQUENCE_STRING_PROJECTION_V2() {
 
 }  // namespace custom
 }  // namespace ops
-}  // namespace tflite
+}  // namespace seq_flow_lite

research/seq_flow_lite/tflite_ops/sequence_string_projection.h

@@ -16,7 +16,7 @@ limitations under the License.
 #define TENSORFLOW_MODELS_SEQUENCE_PROJECTION_TFLITE_OPS_SEQUENCE_STRING_PROJECTION_H_
 #include "tensorflow/lite/kernels/register.h"
 
-namespace tflite {
+namespace seq_flow_lite {
 namespace ops {
 namespace custom {
 
@@ -29,6 +29,6 @@ extern const char kSequenceStringProjectionV2[];
 TfLiteRegistration* Register_SEQUENCE_STRING_PROJECTION_V2();
 }  // namespace custom
 }  // namespace ops
-}  // namespace tflite
+}  // namespace seq_flow_lite
 
 #endif  // TENSORFLOW_MODELS_SEQUENCE_PROJECTION_TFLITE_OPS_SEQUENCE_STRING_PROJECTION_H_

research/seq_flow_lite/tflite_ops/sequence_string_projection_test.cc

@@ -25,29 +25,32 @@ limitations under the License.
 #include "tf_ops/projection_util.h"  // seq_flow_lite
 #include "tflite_ops/tf_tflite_diff_test_util.h"  // seq_flow_lite
 
-namespace tflite {
+namespace seq_flow_lite {
 
 namespace ops {
 namespace custom {
 
 namespace {
 
+using ::seq_flow_lite::testing::AttrValue;
+using ::seq_flow_lite::testing::FloatTensor;
+using ::seq_flow_lite::testing::IntTensor;
+using ::seq_flow_lite::testing::OpEquivTestCase;
+using ::seq_flow_lite::testing::StringTensor;
+using ::seq_flow_lite::testing::TensorflowTfLiteOpTest;
 using ::testing::ElementsAreArray;
-using ::tflite::testing::AttrValue;
-using ::tflite::testing::FloatTensor;
-using ::tflite::testing::IntTensor;
-using ::tflite::testing::OpEquivTestCase;
-using ::tflite::testing::StringTensor;
-using ::tflite::testing::TensorflowTfLiteOpTest;
-
-class SequenceStringProjectionModel : public SingleOpModel {
+using ::tflite::TensorType_FLOAT32;
+using ::tflite::TensorType_STRING;
+using ::tflite::TensorType_UINT8;
+
+class SequenceStringProjectionModel : public ::tflite::SingleOpModel {
  public:
   explicit SequenceStringProjectionModel(
       bool split_on_space, int max_splits, int word_novelty_bits,
-      int doc_size_levels, bool add_eos_tag, TensorType output_type,
+      int doc_size_levels, bool add_eos_tag, ::tflite::TensorType output_type,
       const std::string& token_separators = "",
       bool normalize_repetition = false, float add_first_cap = 0.0,
-      float add_all_caps = 0.0, const string& hashtype = kMurmurHash) {
+      float add_all_caps = 0.0, const std::string& hashtype = kMurmurHash) {
     flexbuffers::Builder fbb;
     fbb.Map([&] {
       fbb.Int("feature_size", 4);
@@ -798,11 +801,11 @@ INSTANTIATE_TEST_SUITE_P(
     SequenceStringProjectionTests, SequenceStringProjectionTest,
     ::testing::ValuesIn(SequenceStringProjectionTestCases()));
 
-class SequenceStringProjectionV2Model : public SingleOpModel {
+class SequenceStringProjectionV2Model : public ::tflite::SingleOpModel {
  public:
   explicit SequenceStringProjectionV2Model(
       std::vector<std::vector<int>> input_shapes,
-      const string& hashtype = kMurmurHash) {
+      const std::string& hashtype = kMurmurHash) {
     flexbuffers::Builder fbb;
     fbb.Map([&] {
       fbb.Int("feature_size", 4);
@@ -827,6 +830,7 @@ class SequenceStringProjectionV2Model : public SingleOpModel {
         << "Cannot allocate tensors";
     return SingleOpModel::InvokeUnchecked();
   }
+  std::vector<int> GetOutputShape() { return GetTensorShape(output_); }
 
  private:
   int input_;
@@ -884,6 +888,15 @@ TEST(SequenceStringProjectionV2Test, RegularInputUint8) {
   m.Invoke({"hello", "world"}, kTfLiteOk);
 }
 
+TEST(SequenceStringProjectionV2Test, NumberProjectionsForMultipleInputs) {
+  SequenceStringProjectionV2Model m({{1, 2}});
+  std::vector<std::string> input = {"hello", "world"};
+  m.Invoke(input, kTfLiteOk);
+  EXPECT_EQ(m.GetOutputShape()[1], input.size());
+  m.Invoke(input, kTfLiteOk);
+  EXPECT_EQ(m.GetOutputShape()[1], input.size());
+}
+
 class SequenceStringProjectionV2Test : public TensorflowTfLiteOpTest {
   std::function<TfLiteRegistration*()> TfLiteOpRegistration() override {
     return ops::custom::Register_SEQUENCE_STRING_PROJECTION_V2;
@@ -986,7 +999,7 @@ INSTANTIATE_TEST_SUITE_P(
 }  // namespace
 }  // namespace custom
 }  // namespace ops
-}  // namespace tflite
+}  // namespace seq_flow_lite
 
 int main(int argc, char** argv) {
   // On Linux, add: absl::SetFlag(&FLAGS_logtostderr, true);

research/seq_flow_lite/tflite_ops/tf_tflite_diff_test_util.cc

@@ -19,11 +19,16 @@ limitations under the License.
 #include "tensorflow/core/framework/op.h"
 #include "tensorflow/core/lib/core/status_test_util.h"
 
-namespace tflite {
+namespace seq_flow_lite {
 namespace testing {
 
 using ::tensorflow::TensorProto;
 using ::testing::FloatNear;
+using ::tflite::TensorType_STRING;
+using ::tflite::TensorType_UINT8;
+using ::tflite::TensorType_INT32;
+using ::tflite::TensorType_BOOL;
+using ::tflite::TensorType_FLOAT32;
 
 ::tflite::TensorType TfTypeToTfLiteType(::tensorflow::DataType dtype) {
   switch (dtype) {
@@ -324,7 +329,7 @@ void TensorflowTfLiteOpTest::CompareOpOutput() {
         const auto& quantization_params =
             GetParam().output_tensors[i].quantization_params;
         if (quantization_params.scale != 0.0) {
-          auto tflite_output_values = Dequantize(
+          auto tflite_output_values = tflite::Dequantize(
               tflite_op_.ExtractVector<uint8_t>(tflite_outputs_[i]),
               quantization_params.scale, quantization_params.zero_point);
           for (int i = 0; i < tf_output_values.size(); i++) {

research/seq_flow_lite/tflite_ops/tf_tflite_diff_test_util.h

@@ -25,7 +25,7 @@ limitations under the License.
 #include "tensorflow/core/kernels/ops_testutil.h"
 #include "tensorflow/lite/kernels/test_util.h"
 
-namespace tflite {
+namespace seq_flow_lite {
 namespace testing {
 
 // Convenience constructors.
@@ -144,6 +144,6 @@ class TensorflowTfLiteOpTest
 };
 
 }  // namespace testing
-}  // namespace tflite
+}  // namespace seq_flow_lite
 
 #endif  // TENSORFLOW_MODELS_SEQUENCE_PROJECTION_TFLITE_OPS_TF_TFLITE_DIFF_TEST_UTIL_H_

research/seq_flow_lite/trainer_v2.py

+# Copyright 2021 The TensorFlow Authors All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+# Lint as: python3
+"""Binary to train PRADO model with TF 2.0."""
+
+import importlib
+import json
+
+from absl import app
+from absl import flags
+from absl import logging
+
+import tensorflow as tf
+
+import input_fn_reader # import root module
+
+FLAGS = flags.FLAGS
+
+flags.DEFINE_string("config_path", None, "Path to a RunnerConfig.")
+flags.DEFINE_enum("runner_mode", "train", ["train", "train_and_eval", "eval"],
+                  "Runner mode.")
+flags.DEFINE_string("master", None, "TensorFlow master URL.")
+flags.DEFINE_string(
+    "output_dir", "/tmp/testV2",
+    "The output directory where the model checkpoints will be written.")
+flags.DEFINE_bool("use_tpu", False, "Whether to use TPU or GPU/CPU.")
+flags.DEFINE_integer(
+    "num_tpu_cores", 8,
+    "Only used if `use_tpu` is True. Total number of TPU cores to use.")
+
+
+def load_runner_config():
+  with tf.io.gfile.GFile(FLAGS.config_path, "r") as f:
+    return json.loads(f.read())
+
+
+def compute_loss(logits, labels, model_config, mode):
+  """Creates a sequence labeling model."""
+  if mode != tf.estimator.ModeKeys.PREDICT:
+    if not model_config["multilabel"]:
+      loss = tf.nn.sparse_softmax_cross_entropy_with_logits(
+          labels=labels, logits=logits)
+    else:
+      loss = tf.nn.sigmoid_cross_entropy_with_logits(
+          labels=labels, logits=logits)
+    loss = tf.reduce_mean(loss)
+  else:
+    loss = None
+
+  return loss
+
+
+def model_fn_builder(runner_config, mode):
+  """Returns `model_fn` closure for TPUEstimator."""
+
+  rel_module_path = "" # empty base dir
+  model = importlib.import_module(rel_module_path + runner_config["name"])
+  model_config = runner_config["model_config"]
+  return model.Encoder(model_config, mode)
+
+
+def main(_):
+  runner_config = load_runner_config()
+
+  if FLAGS.output_dir:
+    tf.io.gfile.makedirs(FLAGS.output_dir)
+
+  train_model = model_fn_builder(runner_config, tf.estimator.ModeKeys.TRAIN)
+  optimizer = tf.keras.optimizers.Adam()
+  train_input_fn = input_fn_reader.create_input_fn(
+      runner_config=runner_config,
+      mode=tf.estimator.ModeKeys.TRAIN,
+      drop_remainder=True)
+  params = {"batch_size": runner_config["batch_size"]}
+  train_ds = train_input_fn(params)
+  train_loss = tf.keras.metrics.Mean(name="train_loss")
+
+  @tf.function
+  def train_step(features):
+    with tf.GradientTape() as tape:
+      logits = train_model(features["projection"], features["seq_length"])
+      loss = compute_loss(logits, features["label"],
+                          runner_config["model_config"],
+                          tf.estimator.ModeKeys.TRAIN)
+    gradients = tape.gradient(loss, train_model.trainable_variables)
+    optimizer.apply_gradients(zip(gradients, train_model.trainable_variables))
+    train_loss(loss)
+
+  for epoch in range(1):
+    train_loss.reset_states()
+    for features in train_ds:
+      train_step(features)
+      step = optimizer.iterations.numpy()
+      if step % 100 == 0:
+        logging.info("Running step %s in epoch %s", step, epoch)
+        logging.info("Training loss: %s, epoch: %s, step: %s",
+                     round(train_loss.result().numpy(), 4), epoch, step)
+
+
+if __name__ == "__main__":
+  app.run(main)