3938 enhance milmodel for torchscript

[yiheng-wang-nv]

Signed-off-by: Yiheng Wang vennw@nvidia.com

Fixes #3938 .

Description

With the changes, MILModel will support torchscript for all mil_mode.

Status

Work in progress

Types of changes

• Non-breaking change (fix or new feature that would not break existing functionality).
• Breaking change (fix or new feature that would cause existing functionality to change).
• New tests added to cover the changes.
• Integration tests passed locally by running ./runtests.sh -f -u --net --coverage.
• Quick tests passed locally by running ./runtests.sh --quick --unittests --disttests.
• In-line docstrings updated.
• Documentation updated, tested make html command in the docs/ folder.

[myron]

this looks good, but shouldn't we support all modes?

[yiheng-wang-nv]

this looks good, but shouldn't we support all modes?

Hi @myron , sure. To support att_trans_pyramid mode, the hook based method may need to be replaced. Let me update this PR for you to take a look.

[yiheng-wang-nv]

Hi @myron , to support att_trans_pyramid mode, the way of using def calc_head seems not feasible. Therefore I defined a new class and put everything of the att_trans_pyramid mode into it. The code may look quite different from your original version, could you please help to review it? Feel free to ask if you have any questions, thanks!

[yiheng-wang-nv]

/black