Skip to content

3432 make vit support torchscript#3782

Merged
wyli merged 5 commits intoProject-MONAI:devfrom
yiheng-wang-nv:3432-make-vit-support-torchscript
Feb 10, 2022
Merged

3432 make vit support torchscript#3782
wyli merged 5 commits intoProject-MONAI:devfrom
yiheng-wang-nv:3432-make-vit-support-torchscript

Conversation

@yiheng-wang-nv
Copy link
Copy Markdown
Contributor

@yiheng-wang-nv yiheng-wang-nv commented Feb 9, 2022 •
edited
Loading

Signed-off-by: Yiheng Wang [email protected]

Fixes #3432 .

Description

This PR is used to fix #3422 , different from the way used in #3433 which requires to remove einops, this PR keeps this library.

During this PR, the direct calling of einops.arrange in forward functions of class SABlock has been replaced by using Rearrange layers in order to support torchscript.

Status

Ready/Work in progress/Hold

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.

@yiheng-wang-nv
make vit support torchscript
2a3d9dd 
Signed-off-by: Yiheng Wang <[email protected]>
wyli
wyli approved these changes Feb 9, 2022
View reviewed changes
Copy link
Copy Markdown
Contributor

@wyli wyli left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

thanks this looks good me, if @ahatamiz has no further comments we merge it today.

@wyli
Copy link
Copy Markdown
Contributor

wyli commented Feb 9, 2022

some typing issue here https://github.com/Project-MONAI/MONAI/runs/5122538609?check_suite_focus=true#step:7:21211

@ericspod
Copy link
Copy Markdown
Member

ericspod commented Feb 9, 2022

There may be some deep-seated incompatibilities between Torchscript and Einops: arogozhnikov/einops#115 The issue with build-gpu/GPU-quick-py3 (PT18+CUDA102) failing may be related, later versions of Pytorch after 1.8 may work, I ran tests with Pytorch 1.10 just now and that version at least works. We may have to state only a certain Pytorch version is scriptable and put the condition on the script test for ViT.

@wyli
Copy link
Copy Markdown
Contributor

wyli commented Feb 10, 2022

/build

@wyli wyli enabled auto-merge (squash) February 10, 2022 08:20
@wyli wyli merged commit 0be3341 into Project-MONAI:dev Feb 10, 2022
wyli pushed a commit to wyli/MONAI that referenced this pull request Feb 10, 2022
@yiheng-wang-nv @wyli
3432 make vit support torchscript (Project-MONAI#3782)
d55e922 
* make vit support torchscript

Signed-off-by: Yiheng Wang <[email protected]>

* add torch version restriction

Signed-off-by: Yiheng Wang <[email protected]>

* change skip decorator order

Signed-off-by: Yiheng Wang <[email protected]>

* remove extra cls

Signed-off-by: Yiheng Wang <[email protected]>
@marksgraham marksgraham mentioned this pull request Nov 21, 2022
Merged
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@wyli wyli wyli approved these changes

@ahatamiz ahatamiz Awaiting requested review from ahatamiz

@ericspod ericspod Awaiting requested review from ericspod ericspod is a code owner

@Nic-Ma Nic-Ma Awaiting requested review from Nic-Ma Nic-Ma is a code owner

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

ViT does not support torchscript

3 participants

@yiheng-wang-nv @wyli @ericspod