How to speed up the transformer inference?
Author: hahadashi
Pull request: no

i meet a problem, if i direct use the model to inference, it is very slow. and i try to split the model (encoder, decoder), i freeze enc ckpt model and dec ckpt model to enc pb model and dec pb model(use tf C++ to run). only compute enc pb model one time and use dec pb model to complete prediction. but its speed is also slow. i kown one reason cause this problem that during i inference the dec model ,there are a lot of repetitive operations.  have your solve the problems。hope you give me some advice or a demo，thx



Comments:
Hi @hahadashi, 

Can you add a code snippet so that we know which model you are using and so that we can reproduce the behavior? 

> Hi @hahadashi,
> 
> Can you add a code snippet so that we know which model you are using and so that we can reproduce the behavior?

thx your response, i use this tutor https://www.tensorflow.org/tutorials/text/transformer#encoder_layer  train the the model

Sorry maybe I was not precise enough:
Which model of the `transformers` library (e.g. Bert, GPT2) did you use? And can you copy / paste the exact code which has a `transformers` model in it that was slow for inference.

> Sorry maybe I was not precise enough:
> Which model of the `transformers` library (e.g. Bert, GPT2) did you use? And can you copy / paste the exact code which has a `transformers` model in it that was slow for inference.

thx your response, 
`class MultiHeadAttention(tf.keras.layers.Layer):
  def __init__(self, d_model, num_heads):
    super(MultiHeadAttention, self).__init__()
    self.num_heads = num_heads
    self.d_model = d_model
    
    assert d_model % self.num_heads == 0
    
    self.depth = d_model // self.num_heads
    
    self.wq = tf.keras.layers.Dense(d_model)
    self.wk = tf.keras.layers.Dense(d_model)
    self.wv = tf.keras.layers.Dense(d_model)
    
    self.dense = tf.keras.layers.Dense(d_model)
        
  def split_heads(self, x, batch_size):
    """分拆最后一个维度到 (num_heads, depth).
    转置结果使得形状为 (batch_size, num_heads, seq_len, depth)
    """
    x = tf.reshape(x, (batch_size, -1, self.num_heads, self.depth))
    return tf.transpose(x, perm=[0, 2, 1, 3])
    
  def call(self, v, k, q, mask):
    batch_size = tf.shape(q)[0]
    
    q = self.wq(q)  # (batch_size, seq_len, d_model)
    k = self.wk(k)  # (batch_size, seq_len, d_model)
    v = self.wv(v)  # (batch_size, seq_len, d_model)
    
    q = self.split_heads(q, batch_size)  # (batch_size, num_heads, seq_len_q, depth)
    k = self.split_heads(k, batch_size)  # (batch_size, num_heads, seq_len_k, depth)
    v = self.split_heads(v, batch_size)  # (batch_size, num_heads, seq_len_v, depth)
    
    # scaled_attention.shape == (batch_size, num_heads, seq_len_q, depth)
    # attention_weights.shape == (batch_size, num_heads, seq_len_q, seq_len_k)
    scaled_attention, attention_weights = scaled_dot_product_attention(
        q, k, v, mask)
    
    scaled_attention = tf.transpose(scaled_attention, perm=[0, 2, 1, 3])  # (batch_size, seq_len_q, num_heads, depth)

    concat_attention = tf.reshape(scaled_attention, 
                                  (batch_size, -1, self.d_model))  # (batch_size, seq_len_q, d_model)

    output = self.dense(concat_attention)  # (batch_size, seq_len_q, d_model)
        
    return output, attention_weights`


during i use the decoder predict, eg, 
1、first i input "sos", get a argmax result "A"
2、 i input "sos A" get a argmax "B"
3、 i input "sos A B" get a argmax "C" 
4......, 
 If we don't save the intermediate state，there are a lot of repetitive operations.
like 
the step 2, sos has compute in step 1, 
step 3, "A" has compute in step 2,
thx

If my idea is wrong, hope you give me some other speed up advices,

It's seems like you're not using the `transformers` repository.

I suggest you use Stack Overflow, where you will more likely receive answers to your question. 

> I suggest you use Stack Overflow, where you will more likely receive answers to your question.

OK, thx