add max_seq_len to DocNN TorchScript model (#1279)

Summary:
Pull Request resolved: #1279

This could dramatically reduce the memory usage for DocNN TorchScript model.

Quick experiment: https://our.intern.facebook.com/intern/anp/view/?id=215724

https://fb.workplace.com/groups/1941258842562334/permalink/3002460646442143/

Reviewed By: m3rlin45

Differential Revision: D20409424

fbshipit-source-id: 1794e2c687c2463b98f1c62cd6842cb1a1b8cda6

Author: chenyangyu1988

pytext/models/doc_model.py

@@ -109,6 +109,14 @@ def torchscriptify(self, tensorizers, traced_model):
         output_layer = self.output_layer.torchscript_predictions()
 
         input_vocab = tensorizers["tokens"].vocab
+        max_seq_len = tensorizers["tokens"].max_seq_len or -1
+
+        """
+        The input tensor packing memory is allocated/cached for different shapes,
+        and max sequence length will help to reduce the number of different tensor
+        shapes. We noticed that the TorchScript model could use 25G for offline
+        inference on CPU without using max_seq_len.
+        """
 
         class Model(jit.ScriptModule):
             def __init__(self):
@@ -117,6 +125,7 @@ def __init__(self):
                 self.model = traced_model
                 self.output_layer = output_layer
                 self.pad_idx = jit.Attribute(input_vocab.idx[PAD], int)
+                self.max_seq_len = jit.Attribute(max_seq_len, int)
 
             @jit.script_method
             def forward(
@@ -128,8 +137,15 @@ def forward(
                 if tokens is None:
                     raise RuntimeError("tokens is required")
 
-                seq_lens = make_sequence_lengths(tokens)
-                word_ids = self.vocab.lookup_indices_2d(tokens)
+                trimmed_tokens: List[List[str]] = []
+                if self.max_seq_len >= 0:
+                    for token in tokens:
+                        trimmed_tokens.append(token[0 : self.max_seq_len])
+                else:
+                    trimmed_tokens = tokens
+
+                seq_lens = make_sequence_lengths(trimmed_tokens)
+                word_ids = self.vocab.lookup_indices_2d(trimmed_tokens)
                 word_ids = pad_2d(word_ids, seq_lens, self.pad_idx)
                 logits = self.model(torch.tensor(word_ids), torch.tensor(seq_lens))
                 return self.output_layer(logits)
@@ -142,6 +158,7 @@ def __init__(self):
                 self.model = traced_model
                 self.output_layer = output_layer
                 self.pad_idx = jit.Attribute(input_vocab.idx[PAD], int)
+                self.max_seq_len = jit.Attribute(max_seq_len, int)
 
             @jit.script_method
             def forward(
@@ -156,8 +173,15 @@ def forward(
                 if dense_feat is None:
                     raise RuntimeError("dense_feat is required")
 
-                seq_lens = make_sequence_lengths(tokens)
-                word_ids = self.vocab.lookup_indices_2d(tokens)
+                trimmed_tokens: List[List[str]] = []
+                if self.max_seq_len >= 0:
+                    for token in tokens:
+                        trimmed_tokens.append(token[0 : self.max_seq_len])
+                else:
+                    trimmed_tokens = tokens
+
+                seq_lens = make_sequence_lengths(trimmed_tokens)
+                word_ids = self.vocab.lookup_indices_2d(trimmed_tokens)
                 word_ids = pad_2d(word_ids, seq_lens, self.pad_idx)
                 dense_feat = self.normalizer.normalize(dense_feat)
                 logits = self.model(