Skip to content

Commit aedf32c

Browse files
meatybobbynv-kkudrynski
meatybobby
authored and
nv-kkudrynski
committed
[BERT/TF2] Add TF-TRT support
1 parent bcc434b commit aedf32c

2 files changed

Lines changed: 98 additions & 14 deletions

File tree

‎TensorFlow2/LanguageModeling/BERT/run_squad.py‎

Lines changed: 28 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -47,21 +47,24 @@
4747
import squad_lib_sp
4848
import tokenization
4949
import gpu_affinity
50+
import tf_trt
5051
from official.utils.misc import distribution_utils
5152
from official.utils.misc import keras_utils
5253
from official.utils.misc import tpu_lib
5354
import dllogger_class
5455

5556
flags.DEFINE_enum(
5657
'mode', 'train_and_predict',
57-
['train_and_predict', 'train', 'predict', 'export_only'],
58-
'One of {"train_and_predict", "train", "predict", "export_only"}. '
58+
['train_and_predict', 'train', 'predict', 'export_only', 'sm_predict', 'trt_predict'],
59+
'One of {"train_and_predict", "train", "predict", "export_only", "sm_predict", "trt_predict"}. '
5960
'`train_and_predict`: both train and predict to a json file. '
6061
'`train`: only trains the model. '
6162
'trains the model and evaluates in the meantime. '
6263
'`predict`: predict answers from the squad json file. '
6364
'`export_only`: will take the latest checkpoint inside '
64-
'model_dir and export a `SavedModel`.')
65+
'model_dir and export a `SavedModel`.'
66+
'`sm_predict`: will load SavedModel from savedmodel_dir and predict answers'
67+
'`trt_predict`: will load SavedModel from savedmodel_dir, convert and predict answers with TF-TRT')
6568
flags.DEFINE_string('train_data_path', '',
6669
'Training data path with train tfrecords.')
6770
flags.DEFINE_string(
@@ -101,6 +104,9 @@
101104
'sp_model_file', None,
102105
'The path to the sentence piece model. Used by sentence piece tokenizer '
103106
'employed by ALBERT.')
107+
flags.DEFINE_string(
108+
'savedmodel_dir', None,
109+
'The path of SavedModel for Savedmodel and TF-TRT prediction.')
104110

105111
common_flags.define_common_bert_flags()
106112

@@ -194,18 +200,25 @@ def predict_squad_customized(strategy, input_meta_data, bert_config,
194200
else:
195201
predict_iterator = iter(predict_dataset_fn())
196202

197-
with distribution_utils.get_strategy_scope(strategy):
198-
squad_model, _ = bert_models.squad_model(
199-
bert_config, input_meta_data['max_seq_length'], float_type=tf.float16 if FLAGS.use_fp16 else tf.float32)
203+
if FLAGS.mode == 'trt_predict':
204+
squad_model = tf_trt.TFTRTModel(FLAGS.savedmodel_dir, "amp" if FLAGS.use_fp16 else "fp32")
200205

201-
if FLAGS.init_checkpoint:
202-
checkpoint = tf.train.Checkpoint(model=squad_model)
203-
checkpoint.restore(FLAGS.init_checkpoint).expect_partial()
206+
elif FLAGS.mode == 'sm_predict':
207+
squad_model = tf_trt.SavedModel(FLAGS.savedmodel_dir, "amp" if FLAGS.use_fp16 else "fp32")
208+
209+
else:
210+
with distribution_utils.get_strategy_scope(strategy):
211+
squad_model, _ = bert_models.squad_model(
212+
bert_config, input_meta_data['max_seq_length'], float_type=tf.float16 if FLAGS.use_fp16 else tf.float32)
213+
214+
if FLAGS.init_checkpoint:
215+
checkpoint = tf.train.Checkpoint(model=squad_model)
216+
checkpoint.restore(FLAGS.init_checkpoint).expect_partial()
204217

205-
checkpoint_path = tf.train.latest_checkpoint(FLAGS.model_dir)
206-
logging.info('Restoring checkpoints from %s', checkpoint_path)
207-
checkpoint = tf.train.Checkpoint(model=squad_model)
208-
checkpoint.restore(checkpoint_path).expect_partial()
218+
checkpoint_path = tf.train.latest_checkpoint(FLAGS.model_dir)
219+
logging.info('Restoring checkpoints from %s', checkpoint_path)
220+
checkpoint = tf.train.Checkpoint(model=squad_model)
221+
checkpoint.restore(checkpoint_path).expect_partial()
209222

210223
@tf.function
211224
def predict_step(iterator):
@@ -621,12 +634,13 @@ def main(_):
621634
policy = tf.keras.mixed_precision.experimental.Policy("mixed_float16")
622635
tf.keras.mixed_precision.experimental.set_policy(policy)
623636

637+
os.makedirs(FLAGS.model_dir, exist_ok=True)
624638
dllogging = dllogger_class.dllogger_class(FLAGS.dllog_path)
625639
input_meta_data['dllogging'] = dllogging
626640

627641
if FLAGS.mode in ('train', 'train_and_predict'):
628642
train_squad(strategy, input_meta_data)
629-
if FLAGS.mode in ('predict', 'train_and_predict') and (not FLAGS.use_horovod or hvd.rank() == 0):
643+
if FLAGS.mode in ('predict', 'sm_predict', 'trt_predict', 'train_and_predict') and (not FLAGS.use_horovod or hvd.rank() == 0):
630644
predict_squad(strategy, input_meta_data)
631645

632646

‎TensorFlow2/LanguageModeling/BERT/tf_trt.py‎

Lines changed: 70 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,70 @@
1+
# Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
# ==============================================================================
15+
import tensorflow as tf
16+
from tensorflow.python.compiler.tensorrt import trt_convert as trt
17+
from tensorflow.compat.v1.saved_model import tag_constants, signature_constants
18+
19+
20+
def export_model(model_dir, prec, tf_trt_model_dir=None):
21+
model = tf.saved_model.load(model_dir)
22+
input_shape = [1, 384]
23+
dummy_input = tf.constant(tf.zeros(input_shape, dtype=tf.int32))
24+
x = [
25+
tf.constant(dummy_input, name='input_word_ids'),
26+
tf.constant(dummy_input, name='input_mask'),
27+
tf.constant(dummy_input, name='input_type_ids'),
28+
]
29+
_ = model(x)
30+
31+
trt_prec = trt.TrtPrecisionMode.FP32 if prec == "fp32" else trt.TrtPrecisionMode.FP16
32+
converter = trt.TrtGraphConverterV2(
33+
input_saved_model_dir=model_dir,
34+
conversion_params=trt.TrtConversionParams(precision_mode=trt_prec),
35+
)
36+
converter.convert()
37+
tf_trt_model_dir = tf_trt_model_dir or f'/tmp/tf-trt_model_{prec}'
38+
converter.save(tf_trt_model_dir)
39+
print(f"TF-TRT model saved at {tf_trt_model_dir}")
40+
41+
class SavedModel:
42+
def __init__(self, model_dir, precision):
43+
self.saved_model_loaded = tf.saved_model.load(model_dir, tags=[tag_constants.SERVING])
44+
self.graph_func = self.saved_model_loaded.signatures[signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY]
45+
self.precision = tf.float16 if precision == "amp" else tf.float32
46+
47+
def __call__(self, x, **kwargs):
48+
return self.infer_step(x)
49+
50+
@tf.function
51+
def infer_step(self, x):
52+
output = self.graph_func(**x)
53+
return output['start_positions'], output['end_positions']
54+
55+
class TFTRTModel:
56+
def __init__(self, model_dir, precision):
57+
temp_tftrt_dir = f"/tmp/tf-trt_model_{precision}"
58+
export_model(model_dir, precision, temp_tftrt_dir)
59+
saved_model_loaded = tf.saved_model.load(temp_tftrt_dir, tags=[tag_constants.SERVING])
60+
print(f"TF-TRT model loaded from {temp_tftrt_dir}")
61+
self.graph_func = saved_model_loaded.signatures[signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY]
62+
self.precision = tf.float16 if precision == "amp" else tf.float32
63+
64+
def __call__(self, x, **kwargs):
65+
return self.infer_step(x)
66+
67+
@tf.function
68+
def infer_step(self, x):
69+
output = self.graph_func(**x)
70+
return output['start_positions'], output['end_positions']

0 commit comments

Comments
 (0)