Skip to content

add classification support for ViT Model #2861

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 5 commits into from
Aug 29, 2021
Merged

add classification support for ViT Model #2861

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
e294f51
add classification support for ViT Model
ahatamiz Aug 29, 2021
c4dfe3b
add classification support for ViT Model
ahatamiz Aug 29, 2021
4e4d133
add classification support for ViT Model
ahatamiz Aug 29, 2021
b9d59b9
Merge branch 'dev' into model_vit_cls_v1
ahatamiz Aug 29, 2021
cd147a0
[MONAI] python code formatting
monai-bot Aug 29, 2021
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion monai/apps/deepedit/transforms.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,10 +6,10 @@

from monai.config import KeysCollection
from monai.transforms.transform import MapTransform, Randomizable, Transform
from monai.utils import optional_import

logger = logging.getLogger(__name__)

from monai.utils import optional_import

distance_transform_cdt, _ = optional_import("scipy.ndimage.morphology", name="distance_transform_cdt")

Expand Down
10 changes: 5 additions & 5 deletions monai/networks/nets/unetr.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -34,9 +34,9 @@ def __init__(
hidden_size: int = 768,
mlp_dim: int = 3072,
num_heads: int = 12,
pos_embed: str = "perceptron",
pos_embed: str = "conv",
norm_name: Union[Tuple, str] = "instance",
conv_block: bool = False,
conv_block: bool = True,
res_block: bool = True,
dropout_rate: float = 0.0,
spatial_dims: int = 3,
Expand All @@ -59,13 +59,13 @@ def __init__(

Examples::

# for single channel input 4-channel output with patch size of (96,96,96), feature size of 32 and batch norm
# for single channel input 4-channel output with image size of (96,96,96), feature size of 32 and batch norm
>>> net = UNETR(in_channels=1, out_channels=4, img_size=(96,96,96), feature_size=32, norm_name='batch')

# for single channel input 4-channel output with patch size of (96,96), feature size of 32 and batch norm
# for single channel input 4-channel output with image size of (96,96), feature size of 32 and batch norm
>>> net = UNETR(in_channels=1, out_channels=4, img_size=96, feature_size=32, norm_name='batch', spatial_dims=2)

# for 4-channel input 3-channel output with patch size of (128,128,128), conv position embedding and instance norm
# for 4-channel input 3-channel output with image size of (128,128,128), conv position embedding and instance norm
>>> net = UNETR(in_channels=4, out_channels=3, img_size=(128,128,128), pos_embed='conv', norm_name='instance')

"""
Expand Down
16 changes: 12 additions & 4 deletions monai/networks/nets/vit.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@

from typing import Sequence, Union

import torch
import torch.nn as nn

from monai.networks.blocks.patchembedding import PatchEmbeddingBlock
Expand All @@ -33,7 +34,7 @@ def __init__(
mlp_dim: int = 3072,
num_layers: int = 12,
num_heads: int = 12,
pos_embed: str = "perceptron",
pos_embed: str = "conv",
classification: bool = False,
num_classes: int = 2,
dropout_rate: float = 0.0,
Expand All @@ -56,12 +57,15 @@ def __init__(

Examples::

# for single channel input with patch size of (96,96,96), conv position embedding and segmentation backbone
# for single channel input with image size of (96,96,96), conv position embedding and segmentation backbone
>>> net = ViT(in_channels=1, img_size=(96,96,96), pos_embed='conv')

# for 3-channel with patch size of (128,128,128), 24 layers and classification backbone
# for 3-channel with image size of (128,128,128), 24 layers and classification backbone
>>> net = ViT(in_channels=3, img_size=(128,128,128), pos_embed='conv', classification=True)

# for 3-channel with image size of (224,224), 12 layers and classification backbone
>>> net = ViT(in_channels=3, img_size=(224,224), pos_embed='conv', classification=True, spatial_dims=2)

"""

super(ViT, self).__init__()
Expand All @@ -88,10 +92,14 @@ def __init__(
)
self.norm = nn.LayerNorm(hidden_size)
if self.classification:
self.classification_head = nn.Linear(hidden_size, num_classes)
self.cls_token = nn.Parameter(torch.zeros(1, 1, hidden_size))
self.classification_head = nn.Sequential(nn.Linear(hidden_size, num_classes), nn.Tanh())

def forward(self, x):
x = self.patch_embedding(x)
if self.classification:
cls_token = self.cls_token.expand(x.shape[0], -1, -1)
x = torch.cat((cls_token, x), dim=1)
hidden_states_out = []
for blk in self.blocks:
x = blk(x)
Expand Down
50 changes: 26 additions & 24 deletions tests/test_vit.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -26,30 +26,32 @@
for num_heads in [12]:
for mlp_dim in [3072]:
for num_layers in [4]:
for num_classes in [2]:
for num_classes in [8]:
for pos_embed in ["conv"]:
# for classification in [False, True]: # TODO: test classification
for nd in (2, 3):
test_case = [
{
"in_channels": in_channels,
"img_size": (img_size,) * nd,
"patch_size": (patch_size,) * nd,
"hidden_size": hidden_size,
"mlp_dim": mlp_dim,
"num_layers": num_layers,
"num_heads": num_heads,
"pos_embed": pos_embed,
"classification": False,
"num_classes": num_classes,
"dropout_rate": dropout_rate,
},
(2, in_channels, *([img_size] * nd)),
(2, (img_size // patch_size) ** nd, hidden_size),
]
if nd == 2:
test_case[0]["spatial_dims"] = 2 # type: ignore
TEST_CASE_Vit.append(test_case)
for classification in [False, True]:
for nd in (2, 3):
test_case = [
{
"in_channels": in_channels,
"img_size": (img_size,) * nd,
"patch_size": (patch_size,) * nd,
"hidden_size": hidden_size,
"mlp_dim": mlp_dim,
"num_layers": num_layers,
"num_heads": num_heads,
"pos_embed": pos_embed,
"classification": classification,
"num_classes": num_classes,
"dropout_rate": dropout_rate,
},
(2, in_channels, *([img_size] * nd)),
(2, (img_size // patch_size) ** nd, hidden_size),
]
if nd == 2:
test_case[0]["spatial_dims"] = 2 # type: ignore
if test_case[0]["classification"]: # type: ignore
test_case[2] = (2, test_case[0]["num_classes"]) # type: ignore
TEST_CASE_Vit.append(test_case)


class TestPatchEmbeddingBlock(unittest.TestCase):
Expand Down Expand Up @@ -113,7 +115,7 @@ def test_ill_arg(self):
num_layers=12,
num_heads=8,
pos_embed="perceptron",
classification=False,
classification=True,
dropout_rate=0.3,
)

Expand Down