Add test for truncated sequence inference (#99)

* Add test for truncated sequence inference with scan model

* Address CR

* Update to Scan opset 9

onnxruntime/test/framework/inference_session_test.cc

@@ -4,11 +4,14 @@
 #include "core/session/inference_session.h"
 
 #include <algorithm>
+#include <cfloat>
 #include <functional>
 #include <iterator>
 #include <thread>
 #include <fstream>
 
+#include <google/protobuf/io/zero_copy_stream_impl.h>
+#include "core/platform/env.h"
 #include "core/common/logging/logging.h"
 #include "core/common/profiler.h"
 #include "core/framework/execution_provider.h"
@@ -1111,5 +1114,158 @@ TEST(ExecutionProviderTest, FunctionInlineTest) {
   VerifyOutputs(fetches, expected_dims_mul_m, expected_values_mul_m);
 }
 
+TEST(InferenceSessionTests, TestTruncatedSequence) {
+  // model/data generated by <repo>/onnxruntime/test/testdata/CNTK/gen.py GenScan()
+  static const std::string LSTM_MODEL_URI = "testdata/scan_1.pb";
+  // This model is a 4x forward LSTM. Parse it to find out mapping between init_state input/output
+  ONNX_NAMESPACE::ModelProto model_proto;
+  int model_fd;
+  auto status = Env::Default().FileOpenRd(LSTM_MODEL_URI, model_fd);
+  ASSERT_TRUE(status.IsOK());
+  google::protobuf::io::FileInputStream f(model_fd);
+  f.SetCloseOnDelete(true);
+  ASSERT_TRUE(model_proto.ParseFromZeroCopyStream(&f));
+  GraphProto& graph_proto = *model_proto.mutable_graph();
+
+  auto find_attr = [&](const NodeProto& node, const std::string& attr_name) -> const AttributeProto* {
+    for (int i = 0; i < node.attribute_size(); ++i) {
+      auto& attr = node.attribute(i);
+      if (attr.name() == attr_name)
+        return &attr;
+    }
+    return nullptr;
+  };
+
+  std::unordered_map<std::string, std::string> init_state_map;
+  for (int i_node = 0; i_node < graph_proto.node_size(); ++i_node) {
+    auto& node = *graph_proto.mutable_node(i_node);
+    if (node.op_type() == "Scan") {
+      // only works in forward, and do not allow bidirection
+      auto attr_directions = find_attr(node, "scan_input_directions");
+      if (attr_directions != nullptr) {
+        ASSERT_TRUE(attr_directions->ints_size() == 1);
+
+        if (attr_directions->ints(0) == 1)
+          continue;  // skip backward Scan
+      }
+
+      // input 0 is optional sequence length, 1..N are for initial states
+      // and N+1..N+num_scan_inputs are actual inputs
+      // output 0..N-1 are for output states, and N.. are actual outputs
+      auto attr_num_scan_inputs = find_attr(node, "num_scan_inputs");
+      ASSERT_TRUE(attr_num_scan_inputs != nullptr);
+      int num_scan_inputs = gsl::narrow_cast<int>(attr_num_scan_inputs->i());
+      ASSERT_TRUE(node.input_size() - num_scan_inputs < node.output_size());
+      for (int i = 0; i < node.input_size() - num_scan_inputs; ++i) {
+        init_state_map.insert(std::make_pair(node.output(i), node.input(i)));
+      }
+    }
+  }
+
+  // now run the truncated model
+  SessionOptions so;
+  InferenceSession session_object(so);
+  ASSERT_TRUE(session_object.Load(LSTM_MODEL_URI).IsOK());
+  ASSERT_TRUE(session_object.Initialize().IsOK());
+
+  RunOptions run_options;
+  run_options.run_tag = "one session/one tag";
+
+  std::vector<int64_t> X_dims = {5, 1, 3};
+  std::vector<float> X = {0.5488135f, 0.71518934f, 0.60276335f,
+                          0.5448832f, 0.4236548f, 0.6458941f,
+                          0.4375872f, 0.891773f, 0.96366274f,
+                          0.3834415f, 0.79172504f, 0.5288949f,
+                          0.56804454f, 0.92559665f, 0.07103606f};
+
+  std::vector<int64_t> Y_dims = {5, 1, 2};
+  std::vector<float> Y_data = {-1.1730184e-04f, -3.1204990e-04f,
+                               -2.9978977e-04f, -1.0602647e-03f,
+                               -3.8115133e-04f, -2.0684483e-03f,
+                               -2.5120965e-04f, -2.9920202e-03f,
+                               3.0980256e-05f, -3.5933927e-03f};
+
+  MLValue ml_value;
+  CreateMLValue<float>(TestCPUExecutionProvider()->GetAllocator(0, OrtMemTypeDefault), X_dims, X, &ml_value);
+
+  std::string input_name = "Input13165";
+  NameMLValMap feeds = {{input_name, ml_value}};
+
+  // prepare outputs for whole sequence
+  std::string final_output_name = "";
+  int final_output_index = -1;
+  for (int i = 0; i < graph_proto.output_size(); ++i) {
+    if (init_state_map.find(graph_proto.output(i).name()) == init_state_map.end()) {
+      ASSERT_TRUE(final_output_name.empty());
+      final_output_name = graph_proto.output(i).name();
+      final_output_index = i;
+    }
+  }
+
+  std::vector<std::string> output_names = {final_output_name};
+  std::vector<MLValue> fetches;
+
+  // Now run the full sequence
+  common::Status st = session_object.Run(run_options, feeds, output_names, &fetches);
+  if (!st.IsOK()) {
+    std::cout << "Run returned status: " << st.ErrorMessage() << std::endl;
+  }
+  ASSERT_TRUE(st.IsOK());
+  ASSERT_EQ(1, fetches.size());
+  auto& rtensor = fetches.front().Get<Tensor>();
+  TensorShape expected_shape(Y_dims);
+  ASSERT_EQ(expected_shape, rtensor.Shape());
+  for (int i = 0; i < Y_data.size(); ++i)
+    EXPECT_NEAR(Y_data[i], rtensor.template Data<float>()[i], FLT_EPSILON);
+
+  // run truncated sequence
+  output_names.clear();
+  for (int i = 0; i < graph_proto.output_size(); ++i) {
+    output_names.push_back(graph_proto.output(i).name());
+  }
+  fetches.clear();
+
+  std::vector<int> truncated_lengths = {2, 2, 1};              // sums to non-truncated length
+  auto seq_stride = TensorShape(X_dims).SizeFromDimension(1);  // sequence is the first dimension of input shape
+  int seq_start = 0;
+  for (auto truncated_len : truncated_lengths) {
+    std::vector<int64_t> truncated_input_dims = X_dims;
+    truncated_input_dims[0] = truncated_len;
+    MLValue truncated_ml_value;
+    std::vector<float> truncated_input(X.begin() + seq_start * seq_stride, X.begin() + (seq_start + truncated_len) * seq_stride);
+    CreateMLValue<float>(TestCPUExecutionProvider()->GetAllocator(0, OrtMemTypeDefault), truncated_input_dims, truncated_input, &truncated_ml_value);
+    NameMLValMap truncated_feeds = {{input_name, truncated_ml_value}};
+    if (seq_start > 0) {
+      // continue from truncated sequence
+      ASSERT_TRUE(fetches.size() == output_names.size());
+      for (int i_output = 0; i_output < output_names.size(); ++i_output) {
+        auto iter = init_state_map.find(output_names[i_output]);
+        if (iter != init_state_map.end())
+          truncated_feeds.insert(std::make_pair(iter->second, fetches[i_output]));
+      }
+    }
+    std::vector<MLValue> truncated_fetches;
+    st = session_object.Run(run_options, truncated_feeds, output_names, &truncated_fetches);
+    if (!st.IsOK()) {
+      std::cout << "Run returned status: " << st.ErrorMessage() << std::endl;
+    }
+    ASSERT_TRUE(st.IsOK());
+
+    // check truncated output
+    auto& truncated_rtensor = truncated_fetches[final_output_index].Get<Tensor>();
+    std::vector<int64_t> truncated_output_dims = Y_dims;
+    truncated_output_dims[0] = truncated_len;
+    TensorShape truncated_shape(truncated_output_dims);
+    ASSERT_EQ(truncated_shape, truncated_rtensor.Shape());
+    auto seq_output_stride = truncated_shape.SizeFromDimension(1);
+    for (int i = 0; i < truncated_shape.Size(); ++i)
+      EXPECT_NEAR(Y_data[i + seq_start * seq_output_stride], truncated_rtensor.template Data<float>()[i], FLT_EPSILON);
+
+    // prepare for next truncated input
+    fetches = truncated_fetches;
+    seq_start += truncated_len;
+  }
+}
+
 }  // namespace test
 }  // namespace onnxruntime

onnxruntime/test/testdata/CNTK/gen.py

@@ -9,40 +9,61 @@ import os
 model_file = 'model.onnx'
 data_dir = 'test_data_set_0'
 
-def SaveTensorProto(file_path, variable, data):
+def SaveTensorProto(file_path, variable, data, name):
     tp = onnx.TensorProto()
-    tp.name = variable.uid
+    tp.name = name if name else variable.uid
+    # ONNX input shape always has sequence axis as the first dimension, if sequence axis exists
     for i in range(len(variable.dynamic_axes)):
-        tp.dims.append(1) # pad 1 for the each dynamic axis
-    for d in variable.shape:
-        tp.dims.append(d)
+        tp.dims.append(data.shape[len(variable.dynamic_axes) - 1 - i])
+    for (i,d) in enumerate(variable.shape):
+        tp.dims.append(d) if d > 0 else tp.dims.append(data.shape[len(data.shape)-len(variable.shape)+i])
     tp.data_type = onnx.TensorProto.FLOAT
     tp.raw_data = data.tobytes()
     with open(file_path, 'wb') as f:
         f.write(tp.SerializeToString())
-        
-def SaveData(test_data_dir, prefix, variables, data_list):
+
+def SaveData(test_data_dir, prefix, variables, data_list, name_replacements=None):
     if isinstance(data_list, np.ndarray):
         data_list = [data_list]
     for (i, d), v in zip(enumerate(data_list), variables):
-        SaveTensorProto(os.path.join(test_data_dir, '{0}_{1}.pb'.format(prefix, i)), v, d)
+        SaveTensorProto(os.path.join(test_data_dir, '{0}_{1}.pb'.format(prefix, i)), v, d, name_replacements[v.uid] if name_replacements else None)
 
-def Save(dir, func, inputs, outputs):
+def Save(dir, func, feed, outputs):
     if not os.path.exists(dir):
         os.makedirs(dir)
-    func.save(os.path.join(dir,model_file), C.ModelFormat.ONNX)
-    
+    onnx_file = os.path.join(dir,model_file)
+    func.save(onnx_file, C.ModelFormat.ONNX)
+
+    # onnx model may have different name for RNN initial states as inputs
+    cntk_to_actual_names = {}
+    with open(onnx_file, 'rb') as ff:
+        sf = ff.read()
+    model = onnx.ModelProto()
+    model.ParseFromString(sf)
+    for actual_input in model.graph.input:
+        actual_input_name = actual_input.name
+        for cntk_input in func.arguments:
+            cntk_name = cntk_input.uid
+            if actual_input_name.startswith(cntk_name):
+                cntk_to_actual_names[cntk_name] = actual_input_name
+
+    if type(feed) is not dict:
+       feed = {func.arguments[0]:feed}
+
+    if type(outputs) is not dict:
+       outputs = {func.outputs[0]:outputs}
+
     test_data_dir = os.path.join(dir, data_dir)
     if not os.path.exists(test_data_dir):
         os.makedirs(test_data_dir)
 
-    SaveData(test_data_dir, 'input', func.arguments, inputs)
-    SaveData(test_data_dir, 'output', func.outputs, outputs)
-        
+    SaveData(test_data_dir, 'input', func.arguments, [feed[var] for var in func.arguments], cntk_to_actual_names)
+    SaveData(test_data_dir, 'output', func.outputs, [outputs[var] for var in func.outputs])
+
 def GenSimple():
     x = C.input_variable((1,3,)) # TODO: fix CNTK exporter bug with shape (3,)
     y = C.layers.Embedding(2)(x) + C.parameter((-1,))
-    data_x = np.random.rand(*x.shape).astype(np.float32)
+    data_x = np.random.rand(1,*x.shape).astype(np.float32)
     data_y = y.eval(data_x)
     Save('test_simple', y, data_x, data_y)
 
@@ -50,10 +71,10 @@ def GenSharedWeights():
     x = C.input_variable((1,3,))
     y = C.layers.Embedding(2)(x)
     y = y + y.parameters[0]
-    data_x = np.random.rand(*x.shape).astype(np.float32)
+    data_x = np.random.rand(1,*x.shape).astype(np.float32)
     data_y = y.eval(data_x)
     Save('test_shared_weights', y, data_x, data_y)
-    
+
 def GenSimpleMNIST():
     input_dim = 784
     num_output_classes = 10
@@ -69,12 +90,111 @@ def GenSimpleMNIST():
 
     model = C.softmax(z)
 
-    data_feature = np.random.rand(*feature.shape).astype(np.float32)
+    data_feature = np.random.rand(1,*feature.shape).astype(np.float32)
     data_output = model.eval(data_feature)
     Save('test_simpleMNIST', model, data_feature, data_output)
 
+def GenMatMul_1k():
+    feature = C.input_variable((1024, 1024,), np.float32)
+    model = C.times(feature, C.parameter((1024,1024), init=C.glorot_uniform()))
+
+    data_feature = np.random.rand(1,*feature.shape).astype(np.float32)
+    data_output = model.eval(data_feature)
+    Save('test_MatMul_1k', model, data_feature, data_output)
+
+def LSTM(cell_dim, use_scan=True):
+    # we now create an LSTM_cell function and call it with the input and placeholders
+    LSTM_cell = C.layers.LSTM(cell_dim)
+
+    @C.Function
+    def func(dh, dc, input):
+        LSTM_func = LSTM_cell(dh, dc, input)
+        if use_scan:
+            LSTM_func_root = C.as_composite(LSTM_func.outputs[0].owner.block_root)
+            args = LSTM_func_root.arguments
+            LSTM_func = LSTM_func_root.clone(C.CloneMethod.share, {args[0]:input, args[1]:dh, args[2]:dc})
+        return LSTM_func
+
+    return func
+
+def GenLSTMx4(use_scan):
+    feature = C.sequence.input_variable((128,), np.float32)
+    lstm1 = C.layers.Recurrence(LSTM(512, use_scan))(feature)
+    lstm2_fw = C.layers.Recurrence(LSTM(512, use_scan))(lstm1)
+    lstm2_bw = C.layers.Recurrence(LSTM(512, use_scan), go_backwards=True)(lstm1)
+    lstm2 = C.splice(lstm2_fw, lstm2_bw, axis=0)
+    lstm3_fw = C.layers.Recurrence(LSTM(512, use_scan))(lstm2)
+    lstm3_bw = C.layers.Recurrence(LSTM(512, use_scan), go_backwards=True)(lstm2)
+    lstm3 = C.splice(lstm3_fw, lstm3_bw, axis=0)
+    lstm4 = C.layers.Recurrence(LSTM(512, use_scan))(lstm3)
+    model = lstm4
+
+    postfix = 'Scan' if use_scan else 'LSTM'
+
+    data_feature = np.random.rand(1,64,128).astype(np.float32)
+    data_output = np.asarray(model.eval(data_feature))
+    Save('test_LSTMx4_' + postfix, model, data_feature, data_output)
+
+def GenScan():
+    np.random.seed(0)
+    feature = C.sequence.input_variable((3,), np.float32)
+    model = C.layers.For(range(4), lambda : C.layers.Recurrence(LSTM(2, use_scan=True)))(feature)
+
+    data_feature = np.random.rand(1,5,3).astype(np.float32)
+    data_output = np.asarray(model.eval(data_feature))
+
+    # print values for test as ground truth
+    print("Scan input\n", data_feature, "\nScan output\n", data_output)
+
+    Save('test_Scan', model, data_feature, data_output)
+
+def GenSimpleScan():
+    feature = C.sequence.input_variable((128,), np.float32)
+    param = C.parameter(shape=(1,), dtype=np.float32)
+    scan = C.layers.Recurrence(lambda h, x: x + h + param)(feature)
+    model = C.sequence.reduce_sum(scan)
+    data_feature = np.random.rand(1,64,128).astype(np.float32)
+    data_output = np.asarray(model.eval(data_feature), dtype=np.float32)
+    Save('test_SimpleScan', model, data_feature, data_output)
+
+def GenLCBLSTM():
+    nc_len = 16
+    nr_len = 16
+    in_dim = 128
+    cell_dim = 256
+    batch_size = 1
+    input = C.sequence.input_variable((in_dim,))
+    init_h = C.input_variable((cell_dim,))
+    init_c = C.input_variable((cell_dim,))
+
+    fwd_cell = LSTM(cell_dim, use_scan=True)
+    fwd_hc = C.layers.RecurrenceFrom(fwd_cell, go_backwards=False, return_full_state=True)(init_h, init_c, input)
+
+    fwd_h_nc = C.sequence.slice(fwd_hc[0], -nr_len-1, -nr_len)
+    fwd_c_nc = C.sequence.slice(fwd_hc[1], -nr_len-1, -nr_len)
+
+    bwd_cell = LSTM(cell_dim, use_scan=True)
+    bwd = C.layers.Recurrence(bwd_cell, go_backwards=True)(input)
+
+    nr = C.splice(fwd_hc[0], bwd)
+    # workaround CNTK bug in slice output name by + 0
+    model = C.combine([C.sequence.reduce_sum(nr), fwd_h_nc + 0, fwd_c_nc + 0])
+    
+    input_data = np.random.rand(batch_size, nc_len+nr_len, in_dim).astype(np.float32)
+    init_h_data = np.random.rand(batch_size, cell_dim).astype(np.float32)
+    init_c_data = np.random.rand(batch_size, cell_dim).astype(np.float32)
+    feed = {input:input_data, init_h:init_h_data, init_c:init_c_data}
+    data_output = model.eval(feed)
+    Save('test_LCBLSTM', model, feed, data_output)
+
 if __name__=='__main__':
     np.random.seed(0)
     GenSimple()
     GenSharedWeights()
     GenSimpleMNIST()
+    GenMatMul_1k()
+    GenLSTMx4(use_scan=True)
+    GenLSTMx4(use_scan=False)
+    GenSimpleScan()
+    GenLCBLSTM()
+    GenScan()