add UNETR, ViT

[ahatamiz]

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

fixes #2403, part of #2394

Description

This pull request adds the following:

• UNETR based on: "Hatamizadeh et al., UNETR: Transformers for 3D Medical Image Segmentation (https://arxiv.org/abs/2103.10504)"
• Vision Transformer (ViT) based on: "Dosovitskiy et al., An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale (https://arxiv.org/abs/2010.11929)"

Status

Ready

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.
• In-line docstrings updated.
• Documentation updated, tested make html command in the docs/ folder.

[ahatamiz]

/black

[ahatamiz]

@wyli , @Nic-Ma

[wyli]

Looks nice! Could you help fix the style errors? https://github.com/Project-MONAI/MONAI/runs/2940634277

[ahatamiz]

Looks nice! Could you help fix the style errors? https://github.com/Project-MONAI/MONAI/runs/2940634277

@wyli, fixed those issues. But there are further issues in lint and type checks which I will address.

[wyli]

thanks, the networks mainly support 3d, would be great to add the 2d support. or we could file a separate feature request and address that later.

[wyli]

might be more robust if there's a torch.flatten as the first step?

[wyli]

Looks nice! Could you help fix the style errors? https://github.com/Project-MONAI/MONAI/runs/2940634277

@wyli, fixed those issues. But there are further issues in lint and type checks which I will address.

sure, in many cases mypy is not smart enough to infer the types and can be difficult to fix... you can use # type: ignore to skip the checks if needed

[wyli]

I think the unit tests took too long to finish on cpu-only testing machine, could you reduce the input size/model size in the unit tests? otherwise please add @skip_if_quick to the test cases, for example

MONAI/tests/test_ahnet.py

Lines 128 to 129 in 12f267c

	@skip_if_quick
	def test_fcn_shape(self, input_param, input_shape, expected_shape):

[ahatamiz]

I think the unit tests took too long to finish on cpu-only testing machine, could you reduce the input size/model size in the unit tests? otherwise please add @skip_if_quick to the test cases, for example

MONAI/tests/test_ahnet.py

Lines 128 to 129 in 12f267c

	@skip_if_quick
	def test_fcn_shape(self, input_param, input_shape, expected_shape):

Pushed another version with reduced input sizes. Tests seem to be done in reasonable amount of time on CPU.

[wyli]

Thanks! It looks nice. Could you please help create a ticket about supporting 2d patch embedding and networks

[ahatamiz]

Thanks! It looks nice. Could you please help create a ticket about supporting 2d patch embedding and networks

Thanks a lot ! I just created #2499 in response to this.