如何让FasterTransformer支持动态batch和动态sequence length

fastertransformer 算子

nvidia在开源的fastertransformer的代码中，提供tensorrt和tensorflow的自定义算子编译和py调用示例，详见fastertransformer.py。但是如果使用tensorflow的自定义算子十分不方便，其batch size 和 sequence length都是固定的。现在提供一种方法让其变成动态的，方法如下：

1. 修改bert_transformer_op.cc，将batch_size，from_seq_len，to_seq_len attr属性去掉，改称input参数，代码如下：

   .input("output_bias: t")
   .input("output_layernorm_beta: t")
   .input("output_layernorm_gamma: t")
+  .input("batch_size: int32")
+  .input("from_seq_len: int32")
   .output("output: t")
   .attr("t: {float, half}")
-  .attr("batch_size: int >= 1")
-  .attr("from_seq_len: int >= 1")
-  .attr("to_seq_len: int >= 1")
+  //.attr("batch_size: int >= 1")
+  //.attr("from_seq_len: int >= 1")
+  //.attr("to_seq_len: int >= 1")
   .attr("head_num: int >= 1")
   .attr("size_per_head: int >= 1")
   .setshapefn([](shape_inference::inferencecontext *c) {
       int batch_size, from_seq_len, to_seq_len, head_num, size_per_head;
-      c->getattr("batch_size", &batch_size);
-      c->getattr("from_seq_len", &from_seq_len);
-      c->getattr("to_seq_len", &to_seq_len);
+      //c->getattr("batch_size", &batch_size);
+      //c->getattr("from_seq_len", &from_seq_len);
+      //c->getattr("to_seq_len", &to_seq_len);
       c->getattr("head_num", &head_num);
       c->getattr("size_per_head", &size_per_head);
-      c->set_output(0, c->makeshape({batch_size * from_seq_len, head_num * size_per_head}));
+      //c->set_output(0, c->makeshape({batch_size * from_seq_len, head_num * size_per_head}));
+      c->set_output(0, c->input(0));
       return status::ok();
       });
 template <typename device, typename t>
@@ -70,14 +71,15 @@ class berttransformerop : public opkernel
   public:
     explicit berttransformerop(opkernelconstruction *context) : opkernel(context)
     {
-      op_requires_ok(context, context->getattr("batch_size", &batch_size_));
-      op_requires_ok(context, context->getattr("from_seq_len", &from_seq_len_));
-      op_requires_ok(context, context->getattr("to_seq_len", &to_seq_len_));
+      //op_requires_ok(context, context->getattr("batch_size", &batch_size_));
+      //op_requires_ok(context, context->getattr("from_seq_len", &from_seq_len_));
+      //op_requires_ok(context, context->getattr("to_seq_len", &to_seq_len_));
       op_requires_ok(context, context->getattr("head_num", &head_num_));
       op_requires_ok(context, context->getattr("size_per_head", &size_per_head_));
 
-      op_requires(context, (from_seq_len_ == to_seq_len_),
-          errors::invalidargument("only support from_seq_len == to_seq_len"));
+      //printf("++++++++ %d =%d \n", from_seq_len_, to_seq_len_)
+      //op_requires(context, (from_seq_len_ == to_seq_len_),
+      ///    errors::invalidargument("only support from_seq_len == to_seq_len"));
 
       try
       {
@@ -95,6 +97,11 @@ class berttransformerop : public opkernel
       bertencodertransformer<encodertraits_> *encoder_transformer_;
       try
       {
+     
+        batch_size_ = context->input(19).flat<int32>().size()/3;
+        from_seq_len_ = context->input(20).flat<int32>().size()/3;
+        to_seq_len_ = from_seq_len_;
+        //printf("==>%d %d\n", batch_size_, from_seq_len_);
         fastertransformer::allocator<allocatortype::tf> allocator_(context);
         encoder_transformer_ = new bertencodertransformer<encodertraits_>(allocator_, 
           batch_size_, from_seq_len_, to_seq_len_, head_num_, size_per_head_);
@@ -104,7 +111,7 @@ class berttransformerop : public opkernel
         op_requires(context, false, errors::internal(error.what()));
       }
       
-      op_requires(context, context->num_inputs() == 19, errors::invalidargument("less input arguments"));
+      op_requires(context, context->num_inputs() == 21, errors::invalidargument("less input arguments"));
 
       encoderinitparam<datatype_> param; //init param here

由于input在cuda的显存中，直接读取input的数值是不可能的(把数值从显存拷贝内存中，比较耗时)，但是我们可以在内存中直接读取形状的size，我们伪造一个形状的size，通过这个size来获取batch_size 和 seq_len。

1. fastertransformer.py修改如下:

    ...
    fast_list_tensor = tf.shape(input_tensor)
    ...
    layer_output = transformer_op_module.bert_transformer(
        layer_input,
        layer_input,
        trainable_vars[0], trainable_vars[2], trainable_vars[4], trainable_vars[1], trainable_vars[3], trainable_vars[5],
        attention_mask,
        trainable_vars[6], trainable_vars[7], trainable_vars[8], trainable_vars[9], trainable_vars[10], trainable_vars[11],
        trainable_vars[12], trainable_vars[13], trainable_vars[14], trainable_vars[15], tf.tile([[1],[2],[3]], [1,fast_list_tensor[0]]),
        tf.tile([[1],[2],[3]], [1,fast_list_tensor[1]]),
        #batch_size=batch_size, 
        #from_seq_len=seq_length, 
        #to_seq_len=seq_length, 
        head_num=num_attention_heads, size_per_head=attention_head_size)

1. 通过以上修改，我们在使用transformer_op_module的时候，就不需要强制指定batch size 和 seq length了， 表示生成模型的时候，类似这么配置：

input_ids = tf.placeholder(tf.int32,(none, none), 'input_ids')
input_mask = tf.placeholder(tf.float32,(none, none), 'input_mask')
input_type_ids = tf.placeholder(tf.int32,(none, none), 'input_type_ids')

便可以生成支持动态batch和动态seq len的tensorflow模型了。