Enable TF-TRT explicit batch mode by tfeher · Pull Request #36379 · tensorflow/tensorflow

tfeher · 2020-01-31T14:59:03Z

This PR adds a new experimental feature to TF-TRT: explicit batch mode. Explicit batch mode is the first step on the way to enable dynamic shapes in TF-TRT (PR #36080).

The main advantage of explicit batch mode is the following: the batch dimension is not a special dimension in this mode, which means a single TRT network can have layers with different batch sizes and more importantly operators can manipulate the batch dimension. A unit test demonstrates this for the squeeze op.

Currently not all TF-TRT ops support explicit batch mode, it is expected to use this feature for elementwise ops. We are soon going to enable all the existing converters to work with this feature.

Currently this feature does not support INT8 calibration.

This feature is not supposed to change any of the existing behavior of TF-TRT in any of the various modes.

The ground work for this PR was already laid down by @pooyadavoodi in PR #34293. The current PR:

updates the TRTEngineOp to use newer TensorRT API that allows explicit batch mode (enqueueV2)
update the converter of squeeze op
add a unit test to demonstrate explicit batch mode using the squeeze op

To turn on this experimental feature, one has to use a rewriter config. Here is an example how to do it:

    root = _model()
    save.save(root, input_saved_model_dir,
              {signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY: root.run})

    rewriter_config_with_trt = rewriter_config_pb2.RewriterConfig()
    rewriter_config_with_trt.optimizers.extend(
        ["constfold", "layout", "constfold"])
    rewriter_config_with_trt.meta_optimizer_iterations = (
        rewriter_config_pb2.RewriterConfig.ONE)
    optimizer = rewriter_config_with_trt.custom_optimizers.add()
    rewriter_config_with_trt.custom_optimizers.add().name = "constfold"
    optimizer.name = "TensorRTOptimizer"
    optimizer.parameter_map["minimum_segment_size"].i = 1
    optimizer.parameter_map["is_dynamic_op"].b = True
    optimizer.parameter_map["maximum_cached_engines"].i = 1
    optimizer.parameter_map["use_implicit_batch"].b = False

    conversion_params = trt_convert.DEFAULT_TRT_CONVERSION_PARAMS._replace(
        rewriter_config_template=rewriter_config_with_trt)
    converter = trt_convert.TrtGraphConverterV2(input_saved_model_dir=input_saved_model_dir,
                                    conversion_params=conversion_params)
    converter.convert()

    output_saved_model_dir = self.mkdtemp()
    converter.save(output_saved_model_dir=output_saved_model_dir)

Tagging @bixia1 for review.

tfeher · 2020-02-03T15:46:28Z

Update: this message was intended for the discussion of #36080, I have also added this comment there.

Here is a list of PRs related to ~~the current one~~ #36080. They incrementally add changes that shall make #36080 smaller. Each PR is based on the previous one in the list, please review them in the given order:
1 #36379
2. #36434
3. #36435
4. #36439
5. Add optimizaton profiles (Work in progress)
6. Handle multiple contexts (work in progress)
7. Test deserialization of opt profiles (work in progress)

After these ~~the current~~ #36080 will be reduced to the following: Enable build mode to generate optimization profiles for dynamic shapes.

bixia1

Thanks for the work! Looks good to me, I only have some minor comments.

tfeher

Thanks @bixia1 for the review. I have addressed all the issues please have a look.

bixia1

I only see some minor formatting issue, which will hopefully be automatically fixed when the code is imported to google3.

PiperOrigin-RevId: 293412014 Change-Id: I5a0852d415b149e2e39a2e41eb953ad1c8b911d2

rthadur · 2020-02-06T17:14:53Z

Seems auto-merge is not happening but the changes are now committed so we can close this. Thank you for the PR.

tfeher added 8 commits January 31, 2020 13:33

Enable TRTEngine execution in explicit batch mode

eaba0ce

Add DebugString functions for vector of shapes into TensorRT utils

0091af5

Keep logger alive to avoid segfault in trt_engine_resource_ops test

faa3abc

Update doc and error messages for TRT converters build mode

0183511

Add arg to set implicit_batch parameter for GetRewriterConfig

89b53f2

Enable squeeze op for explicit batch dim

15a0b37

Add python test for TF-TRT explicit batch mode

2d98880

Update description of explicit batch test

bc45866

tensorflow-bot Bot added the size:L CL Change Size: Large label Jan 31, 2020

googlebot added the cla: yes label Jan 31, 2020

tfeher mentioned this pull request Jan 31, 2020

Enable explicit batch mode with optimization profiles #36080

Closed

bixia1 self-requested a review January 31, 2020 16:17

rthadur self-assigned this Jan 31, 2020

This was referenced Feb 3, 2020

Improve TensorRT binding index query #36434

Merged

Add TensorRT binding size dimension specification #36435

Merged

bixia1 reviewed Feb 3, 2020

View reviewed changes

tfeher added 10 commits February 4, 2020 15:48

Fix comments and use absl

10d709a

Document SqueezeTensor

fdf52bd

Simplify DebugString functions

86c3de4

Improve error checks

fe8e289

Use correct year in copyright statement

18f1e3f

Merge branch 'master' into trt_explicit_batch_mode

22bf87a

Rename explicit_batch_test.py to trt_mode_test.py

662f9cb

Improve class hierarchy in trt_mode_test.py

50f9a81

Set build_called_once to true after we call build mode

47ad310

Do not run base class for trt mode tests

9c676bd

tfeher commented Feb 4, 2020

View reviewed changes

rthadur requested a review from bixia1 February 4, 2020 17:55

bixia1 reviewed Feb 4, 2020

View reviewed changes

Comment thread tensorflow/compiler/tf2tensorrt/convert/convert_nodes.cc Resolved

Comment thread tensorflow/compiler/tf2tensorrt/convert/convert_nodes.h Resolved

Comment thread tensorflow/python/compiler/tensorrt/test/trt_mode_test.py Resolved

Edit docstring for SqueezeTensor

2024a0d

bixia1 previously approved these changes Feb 4, 2020

View reviewed changes

tensorflow-bot Bot added kokoro:force-run Tests on submitted change ready to pull PR ready for merge process labels Feb 4, 2020

kokoro-team removed the kokoro:force-run Tests on submitted change label Feb 4, 2020

Fix linter errors int trt_mode_test

83d553c

tfeher dismissed bixia1’s stale review via 83d553c February 5, 2020 11:00

tensorflow-bot Bot removed the ready to pull PR ready for merge process label Feb 5, 2020

bixia1 approved these changes Feb 5, 2020

View reviewed changes

tensorflow-bot Bot added kokoro:force-run Tests on submitted change ready to pull PR ready for merge process labels Feb 5, 2020

kokoro-team removed the kokoro:force-run Tests on submitted change label Feb 5, 2020

tensorflow-copybara pushed a commit that referenced this pull request Feb 5, 2020

Merge pull request #36379 from tfeher:trt_explicit_batch_mode

950d9b9

PiperOrigin-RevId: 293412014 Change-Id: I5a0852d415b149e2e39a2e41eb953ad1c8b911d2

rthadur closed this Feb 6, 2020

tfeher mentioned this pull request Feb 10, 2020

Define TensorRT network with dynamic shapes #36439

Merged

tfeher mentioned this pull request Apr 1, 2020

Refactor ExecuteTrtEngine. #38118

Merged

tfeher mentioned this pull request Dec 8, 2020

TF-TRT Dynamic Shapes Feature Tracker #45481

Closed

tfeher deleted the trt_explicit_batch_mode branch May 16, 2022 08:45

Conversation

tfeher commented Jan 31, 2020 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

tfeher commented Feb 3, 2020 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

bixia1 left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

tfeher left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

bixia1 left a comment

Choose a reason for hiding this comment

Uh oh!

rthadur commented Feb 6, 2020

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

tfeher commented Jan 31, 2020 •

edited

Loading

tfeher commented Feb 3, 2020 •

edited

Loading