add swin_unetr model #4074

ahatamiz · 2022-04-04T22:35:38Z

Signed-off-by: ahatamizadeh ahatamizadeh@nvidia.com

Fixes #3520 .

Description

This PR adds the Swin UNETR [1] model to MONAI.

[1]: Hatamizadeh, Ali, Vishwesh Nath, Yucheng Tang, Dong Yang, Holger Roth, and Daguang Xu. "Swin UNETR: Swin Transformers for Semantic Segmentation of Brain Tumors in MRI Images." arXiv preprint arXiv:2201.01266 (2022).

Swin UNETR is also used in self-supervised learning for 3D medical image segmentation [2].

[2]: Tang, Yucheng, Dong Yang, Wenqi Li, Holger Roth, Bennett Landman, Daguang Xu, Vishwesh Nath, and Ali Hatamizadeh. "Self-supervised pre-training of swin transformers for 3d medical image analysis." arXiv preprint arXiv:2111.14791 (2021).

Status

Ready

Types of changes

Non-breaking change (fix or new feature that would not break existing functionality).
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.

Signed-off-by: ahatamizadeh <ahatamizadeh@nvidia.com>

ahatamiz · 2022-04-04T23:04:20Z

/build

ahatamiz · 2022-04-04T23:04:53Z

/black

Signed-off-by: ahatamizadeh <ahatamizadeh@nvidia.com>

ahatamiz · 2022-04-04T23:40:41Z

/build

ahatamiz · 2022-04-04T23:40:54Z

/black

Signed-off-by: ahatamizadeh <ahatamizadeh@nvidia.com>

ahatamiz · 2022-04-05T01:02:27Z

/build

ahatamiz · 2022-04-05T01:02:34Z

/black

Signed-off-by: ahatamizadeh <ahatamizadeh@nvidia.com>

ahatamiz · 2022-04-05T03:08:33Z

/build

ahatamiz · 2022-04-05T03:08:44Z

/black

Signed-off-by: ahatamizadeh <ahatamizadeh@nvidia.com>

ahatamiz · 2022-04-05T05:05:35Z

/build

ahatamiz · 2022-04-05T05:05:42Z

/black

Signed-off-by: ahatamizadeh <ahatamizadeh@nvidia.com>

ahatamiz · 2022-05-02T18:40:39Z

/black

ahatamiz · 2022-05-02T18:41:01Z

/build

Signed-off-by: ahatamizadeh <ahatamizadeh@nvidia.com>

ahatamiz · 2022-05-02T20:40:55Z

/black

ahatamiz · 2022-05-02T20:41:05Z

/build

Signed-off-by: ahatamizadeh <ahatamizadeh@nvidia.com>

ahatamiz · 2022-05-03T01:52:41Z

/black

ahatamiz · 2022-05-03T05:18:04Z

All checks are passed (blossom needs to be triggered) . It would be great if this PR can be reviewed.

Thanks

Nic-Ma · 2022-05-03T05:40:09Z

/build

Nic-Ma · 2022-05-03T05:48:05Z

@wyli , please see the changes as listed below:

2D support added. window_reverse, WindowAttention, SwinTransformerBlock and others support both 2D and 3D. Furthermore, additional 2D test cases were added.

DropPath moved to layers. Test cases are added.

Weigh_init moved to layers. Test cases are added.

PatchEmbed moved under patchembedding.py. Descriptions added for the new embedding (differences etc.)

torch meshgrid index added.

I investigated the possibility of using torch.jit.script. But window operations as well as gradient checkpointing are not supported. Since the input shape needs to be fixed for this model, we can still support Torch Script by using torch.jit.trace ( similar to the approach used in the latest Clara Train release).

Thanks

Hi @wyli ,

Seems @ahatamiz addressed several of your previous comments, could you please help review and confirm them?

Thanks in advance.

ahatamiz

All concerns addressed.

wyli · 2022-05-03T19:52:23Z

/build

wyli

thanks, it looks good to me apart from the typos (commented inline).

optionally the norm_layer could be replaced using get_norm_layer, it supports keyword arguments:

>>> from monai.networks.layers import get_norm_layer
# layernorm with additional params.
>>> norm_name = ("layer", {"normalized_shape": 10})
>>> get_norm_layer(norm_name)
LayerNorm((10,), eps=1e-05, elementwise_affine=True)
# group norm with additional params.
>>> norm_name = ("group", {"num_groups": 1, "num_channels": 10})
>>> get_norm_layer(norm_name)
GroupNorm(1, 10, eps=1e-05, affine=True)
>>>

monai/networks/layers/weight_init.py

monai/networks/nets/swin_unetr.py

Signed-off-by: ahatamizadeh <ahatamizadeh@nvidia.com>

for more information, see https://pre-commit.ci

ahatamiz · 2022-05-03T21:43:23Z

/black

Signed-off-by: monai-bot <monai.miccai2019@gmail.com>

wyli · 2022-05-03T22:32:41Z

/build

* add swin_unetr model (#4074) * add swin_unetr model * 4217 Update PyTorch docker to 22.04 (#4218) * [DLMED] update to 22.04 Signed-off-by: Nic Ma <nma@nvidia.com> * fixes unit test tests.test_lr_finder Signed-off-by: Wenqi Li <wenqil@nvidia.com> * test new_empty Signed-off-by: Wenqi Li <wenqil@nvidia.com> Co-authored-by: Wenqi Li <wenqil@nvidia.com> Co-authored-by: Wenqi Li <831580+wyli@users.noreply.github.com> * Add InstanceNorm3dNVFuser support (#4194) * implement the base class Signed-off-by: Yiheng Wang <vennw@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * add unittest Signed-off-by: Yiheng Wang <vennw@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * autofix Signed-off-by: Yiheng Wang <vennw@nvidia.com> * switch to call apex directly Signed-off-by: Yiheng Wang <vennw@nvidia.com> * uncomment unittest Signed-off-by: Yiheng Wang <vennw@nvidia.com> * add apex install link in docstring Signed-off-by: Yiheng Wang <vennw@nvidia.com> * add channels_last_3d test case Signed-off-by: Yiheng Wang <vennw@nvidia.com> * rewrite types Signed-off-by: Yiheng Wang <vennw@nvidia.com> * change types Signed-off-by: Yiheng Wang <vennw@nvidia.com> * add docstrings Signed-off-by: Yiheng Wang <vennw@nvidia.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Wenqi Li <831580+wyli@users.noreply.github.com> Co-authored-by: Ali Hatamizadeh <ahatamizadeh@nvidia.com> Co-authored-by: Nic Ma <nma@nvidia.com> Co-authored-by: Yiheng Wang <68361391+yiheng-wang-nv@users.noreply.github.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

* add swin_unetr model

* add swin_unetr model (#4074) * add swin_unetr model Signed-off-by: kbressem <kenobressem@gmail.com> * 4217 Update PyTorch docker to 22.04 (#4218) * [DLMED] update to 22.04 Signed-off-by: Nic Ma <nma@nvidia.com> * fixes unit test tests.test_lr_finder Signed-off-by: Wenqi Li <wenqil@nvidia.com> * test new_empty Signed-off-by: Wenqi Li <wenqil@nvidia.com> Co-authored-by: Wenqi Li <wenqil@nvidia.com> Co-authored-by: Wenqi Li <831580+wyli@users.noreply.github.com> Signed-off-by: kbressem <kenobressem@gmail.com> * Add InstanceNorm3dNVFuser support (#4194) * implement the base class Signed-off-by: Yiheng Wang <vennw@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * add unittest Signed-off-by: Yiheng Wang <vennw@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * autofix Signed-off-by: Yiheng Wang <vennw@nvidia.com> * switch to call apex directly Signed-off-by: Yiheng Wang <vennw@nvidia.com> * uncomment unittest Signed-off-by: Yiheng Wang <vennw@nvidia.com> * add apex install link in docstring Signed-off-by: Yiheng Wang <vennw@nvidia.com> * add channels_last_3d test case Signed-off-by: Yiheng Wang <vennw@nvidia.com> * rewrite types Signed-off-by: Yiheng Wang <vennw@nvidia.com> * change types Signed-off-by: Yiheng Wang <vennw@nvidia.com> * add docstrings Signed-off-by: Yiheng Wang <vennw@nvidia.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Wenqi Li <831580+wyli@users.noreply.github.com> Signed-off-by: kbressem <kenobressem@gmail.com> * Update dice.py (#4234) * Update dice.py reduce redundant operations in DiceFocalLoss, initially caused oom Signed-off-by: Ryan Clanton <55164720+ryancinsight@users.noreply.github.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci Signed-off-by: Ryan Clanton <55164720+ryancinsight@users.noreply.github.com> * [MONAI] python code formatting Signed-off-by: monai-bot <monai.miccai2019@gmail.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: monai-bot <monai.miccai2019@gmail.com> Signed-off-by: kbressem <kenobressem@gmail.com> * Bug fix and improvement in WSI (#4216) * Make all transforms optional Signed-off-by: Behrooz <3968947+drbeh@users.noreply.github.com> * Update wsireader tests Signed-off-by: Behrooz <3968947+drbeh@users.noreply.github.com> * Remove optional from PersistentDataset and its derivatives Signed-off-by: Behrooz <3968947+drbeh@users.noreply.github.com> * Add unittests for cache without transform Signed-off-by: Behrooz <3968947+drbeh@users.noreply.github.com> * Add default replace_rate Signed-off-by: Behrooz <3968947+drbeh@users.noreply.github.com> * Add default value Signed-off-by: Behrooz <3968947+drbeh@users.noreply.github.com> * Set default replace_rate to 0.1 Signed-off-by: Behrooz <3968947+drbeh@users.noreply.github.com> * Update metadata to include path Signed-off-by: Behrooz <3968947+drbeh@users.noreply.github.com> * Adds SmartCachePatchWSIDataset Signed-off-by: Behrooz <3968947+drbeh@users.noreply.github.com> * Add unittests for SmartCachePatchWSIDataset Signed-off-by: Behrooz <3968947+drbeh@users.noreply.github.com> * Update references Signed-off-by: Behrooz <3968947+drbeh@users.noreply.github.com> * Update docs Signed-off-by: Behrooz <3968947+drbeh@users.noreply.github.com> * Remove smart cache Signed-off-by: Behrooz <3968947+drbeh@users.noreply.github.com> * Remove unused imports Signed-off-by: Behrooz <3968947+drbeh@users.noreply.github.com> * Add path metadata for OpenSlide Signed-off-by: Behrooz <3968947+drbeh@users.noreply.github.com> * Update metadata to be unified across different backends Signed-off-by: Behrooz <3968947+drbeh@users.noreply.github.com> * Update wsi metadata for multi wsi objects Signed-off-by: Behrooz <3968947+drbeh@users.noreply.github.com> * Add unittests for wsi metadata Signed-off-by: Behrooz <3968947+drbeh@users.noreply.github.com> Signed-off-by: kbressem <kenobressem@gmail.com> * Replace module (#4245) * replace modules Signed-off-by: Richard Brown <33289025+rijobro@users.noreply.github.com> * fix Signed-off-by: Richard Brown <33289025+rijobro@users.noreply.github.com> * replace_module -> replace_modules Signed-off-by: Richard Brown <33289025+rijobro@users.noreply.github.com> * fix Signed-off-by: Richard Brown <33289025+rijobro@users.noreply.github.com> Signed-off-by: kbressem <kenobressem@gmail.com> * Add GaussianSmooth as antialiasing filter in Resize (#4249) Signed-off-by: Can Zhao <canz@nvidia.com> Signed-off-by: kbressem <kenobressem@gmail.com> * 4235 fix 2204 nvfuser issue (#4241) * reproduce issue Signed-off-by: Yiheng Wang <vennw@nvidia.com> * remove 22.01 02 Signed-off-by: Yiheng Wang <vennw@nvidia.com> * remove other workflows Signed-off-by: Yiheng Wang <vennw@nvidia.com> * run on pull request Signed-off-by: Yiheng Wang <vennw@nvidia.com> * remove sleep Signed-off-by: Yiheng Wang <vennw@nvidia.com> * test single layer forward Signed-off-by: Yiheng Wang <vennw@nvidia.com> * add has_nvfuser Signed-off-by: Yiheng Wang <vennw@nvidia.com> * add check within factory Signed-off-by: Yiheng Wang <vennw@nvidia.com> * revert to original cron.yml Signed-off-by: Yiheng Wang <vennw@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix old pt issue Signed-off-by: Yiheng Wang <vennw@nvidia.com> * change to return directly if no cuda Signed-off-by: Yiheng Wang <vennw@nvidia.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Signed-off-by: kbressem <kenobressem@gmail.com> * Update to Bundle Specifiation (#4250) * Update to bundle specifiation Signed-off-by: Eric Kerfoot <eric.kerfoot@kcl.ac.uk> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Adding description in spec discussing the saved Torchscript object's file storage behaviour, and tweaking ckpt_export to add .json extension Signed-off-by: Eric Kerfoot <eric.kerfoot@kcl.ac.uk> * Annotating optional bundle files Signed-off-by: Eric Kerfoot <eric.kerfoot@kcl.ac.uk> * Adjusted ckpt_export test Signed-off-by: Eric Kerfoot <eric.kerfoot@kcl.ac.uk> * Fix Signed-off-by: Eric Kerfoot <eric.kerfoot@kcl.ac.uk> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Signed-off-by: kbressem <kenobressem@gmail.com> * Implement NrrdReader and NrrdImage classes Signed-off-by: kbressem <kenobressem@gmail.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci Signed-off-by: kbressem <kenobressem@gmail.com> * run auto style fixes on image_reader.py Signed-off-by: kbressem <kenobressem@gmail.com> * add NrrdReader to monai/data/__init__.py Signed-off-by: kbressem <kenobressem@gmail.com> * Change the way spatial information is handled in NrrdReader Signed-off-by: kbressem <kenobressem@gmail.com> * add tests for NrrdReader Signed-off-by: kbressem <kenobressem@gmail.com> * Add NrrdReader to list of possible readers for LoadImage Signed-off-by: kbressem <kenobressem@gmail.com> * autofix formating Signed-off-by: kbressem <kenobressem@gmail.com> * autofix formating Signed-off-by: kbressem <kenobressem@gmail.com> * change NrrdImage class to namedtuple and make flake8 happy Signed-off-by: kbressem <kenobressem@gmail.com> * Add pynrrd to requirements Signed-off-by: kbressem <kenobressem@gmail.com> * correct typing for namedtumple make flake8 happy Signed-off-by: kbressem <kenobressem@gmail.com> * Add pynrrd info to `get_optional_config_values` Changed NrrdImage to dataclass Signed-off-by: kbressem <kenobressem@gmail.com> * exclude test_nrrd_reader.py from min tests Signed-off-by: kbressem <kenobressem@gmail.com> * add pynrrd to config files Signed-off-by: kbressem <kenobressem@gmail.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Change the way space is handled in the header. Now, if space is not in header, it is assumed to be LPS and converted to RAS. If space is defined and not LPS, nothing is done to prevent wrong conversions. Signed-off-by: kbressem <kenobressem@gmail.com> * add `TestLoadSaveNrrd` where it is tested if a nrrd file, created by ITKWriter can be loaded again. 2D and 3D files with no channels are tested Signed-off-by: kbressem <kenobressem@gmail.com> * autofix format Co-authored-by: kbressem <kenobressem@charite.de>

ahatamiz added 2 commits April 4, 2022 15:33

add swin_unetr model

f97cad8

Signed-off-by: ahatamizadeh <ahatamizadeh@nvidia.com>

add swin_unetr model

34a1f9a

Signed-off-by: ahatamizadeh <ahatamizadeh@nvidia.com>

ahatamiz force-pushed the swin_unetr_v1 branch from 5151f8b to 34a1f9a Compare April 4, 2022 23:02

ahatamiz requested review from Nic-Ma and wyli April 4, 2022 23:03

add swin_unetr model

c292d2c

Signed-off-by: ahatamizadeh <ahatamizadeh@nvidia.com>

ahatamiz force-pushed the swin_unetr_v1 branch from 9cb10ed to c292d2c Compare April 4, 2022 23:08

add swin_unetr model

83c43de

Signed-off-by: ahatamizadeh <ahatamizadeh@nvidia.com>

ahatamiz force-pushed the swin_unetr_v1 branch from 1f4cbf6 to 83c43de Compare April 4, 2022 23:40

add swin_unetr model

384d1e5

Signed-off-by: ahatamizadeh <ahatamizadeh@nvidia.com>

ahatamiz force-pushed the swin_unetr_v1 branch from c23859e to 384d1e5 Compare April 5, 2022 01:01

add swin_unetr model

1774610

Signed-off-by: ahatamizadeh <ahatamizadeh@nvidia.com>

ahatamiz force-pushed the swin_unetr_v1 branch from 10b3016 to 1774610 Compare April 5, 2022 03:07

add swin_unetr model

a6dec2d

Signed-off-by: ahatamizadeh <ahatamizadeh@nvidia.com>

ahatamiz force-pushed the swin_unetr_v1 branch from 2bccf19 to a6dec2d Compare April 5, 2022 05:04

ahatamiz added this to the Swin UNETR [P1] milestone Apr 5, 2022

ahatamiz self-assigned this Apr 5, 2022

ahatamiz added the Feature request label Apr 5, 2022

ahatamiz force-pushed the swin_unetr_v1 branch from c21afc0 to a6dec2d Compare April 5, 2022 05:41

add swin_unetr model

ea2e6aa

Signed-off-by: ahatamizadeh <ahatamizadeh@nvidia.com>

add swin_unetr model

2631fb3

Signed-off-by: ahatamizadeh <ahatamizadeh@nvidia.com>

ahatamiz force-pushed the swin_unetr_v1 branch from 566d27a to 2631fb3 Compare May 2, 2022 18:39

add swin_unetr model

d1e0817

Signed-off-by: ahatamizadeh <ahatamizadeh@nvidia.com>

ahatamiz force-pushed the swin_unetr_v1 branch from 1d5466d to d1e0817 Compare May 2, 2022 20:39

add swin_unetr model

b6795e4

Signed-off-by: ahatamizadeh <ahatamizadeh@nvidia.com>

ahatamiz force-pushed the swin_unetr_v1 branch from af02e6d to b6795e4 Compare May 3, 2022 01:51

ahatamiz commented May 3, 2022

View reviewed changes

wyli approved these changes May 3, 2022

View reviewed changes

monai/networks/layers/weight_init.py Outdated Show resolved Hide resolved

monai/networks/nets/swin_unetr.py Outdated Show resolved Hide resolved

add swin_unetr model

14f1453

Signed-off-by: ahatamizadeh <ahatamizadeh@nvidia.com>

ahatamiz force-pushed the swin_unetr_v1 branch from 7baaa67 to 14f1453 Compare May 3, 2022 21:18

ahatamiz and others added 2 commits May 3, 2022 14:18

Merge branch 'dev' into swin_unetr_v1

66a4d69

[pre-commit.ci] auto fixes from pre-commit.com hooks

f5c64e1

for more information, see https://pre-commit.ci

[MONAI] python code formatting

ee8ca91

Signed-off-by: monai-bot <monai.miccai2019@gmail.com>

wyli enabled auto-merge (squash) May 3, 2022 22:33

wyli merged commit e093d0b into Project-MONAI:dev May 3, 2022

ahatamiz deleted the swin_unetr_v1 branch May 3, 2022 23:27

Can-Zhao pushed a commit to Can-Zhao/MONAI that referenced this pull request May 10, 2022

add swin_unetr model (Project-MONAI#4074)

48f8b6b

* add swin_unetr model

add swin_unetr model #4074

add swin_unetr model #4074

Uh oh!

Conversation

ahatamiz commented Apr 4, 2022 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Status

Types of changes

Uh oh!

ahatamiz commented Apr 4, 2022

Uh oh!

ahatamiz commented Apr 4, 2022

Uh oh!

ahatamiz commented Apr 4, 2022

Uh oh!

ahatamiz commented Apr 4, 2022

Uh oh!

ahatamiz commented Apr 5, 2022

Uh oh!

ahatamiz commented Apr 5, 2022

Uh oh!

ahatamiz commented Apr 5, 2022

Uh oh!

ahatamiz commented Apr 5, 2022

Uh oh!

ahatamiz commented Apr 5, 2022

Uh oh!

ahatamiz commented Apr 5, 2022

Uh oh!

ahatamiz commented May 2, 2022

Uh oh!

ahatamiz commented May 2, 2022

Uh oh!

ahatamiz commented May 2, 2022

Uh oh!

ahatamiz commented May 2, 2022

Uh oh!

ahatamiz commented May 3, 2022

Uh oh!

ahatamiz commented May 3, 2022

Uh oh!

Nic-Ma commented May 3, 2022

Uh oh!

Nic-Ma commented May 3, 2022

Uh oh!

ahatamiz left a comment

Choose a reason for hiding this comment

Uh oh!

wyli commented May 3, 2022

Uh oh!

wyli left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

ahatamiz commented May 3, 2022

Uh oh!

wyli commented May 3, 2022

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

ahatamiz commented Apr 4, 2022 •

edited

Loading