Repository: Project-MONAI/MONAI

Issue:
Add support for Transformer Optimization Scheduler
**Is your feature request related to a problem? Please describe.**
Add schedulers for training transformer-based models (e.g. ViT). These schedulers include:

- WarmupConstantSchedule
- WarmupLinearSchedule
- WarmupCosineSchedule
- WarmupCosineWithHardRestartsSchedule

**Describe the solution you'd like**
Currently, Pytorch does not have implementations for these schedulers. However, HuggingFace library provides such implementations:
https://huggingface.co/transformers/v1.2.0/_modules/pytorch_transformers/optimization.html


enhance unetr docstring with examples
### Discussed in https://github.com/Project-MONAI/MONAI/discussions/2547

<div type='discussions-op-text'>

<sup>Originally posted by **jhkang0526** July  7, 2021</sup>
hello,

I recently saw about the UNETR. 

Is it currently possible to apply to 3D data?

If possible, how should I use it?

Thanks,

</div>


Discussion:
@wyli @Nic-Ma 
In particular [WarmupCosineSchedule](https://huggingface.co/transformers/v1.2.0/_modules/pytorch_transformers/optimization.html#WarmupCosineSchedule) is widely used for scheduling the optimization in training of Transformers. This scheduler applies linear warmup and then cosine decay. The closest Pytorch implementation is [CosineAnnealingWarmRestarts](https://pytorch.org/docs/stable/_modules/torch/optim/lr_scheduler.html#CosineAnnealingWarmRestarts) which is quite different as it does not apply the linear warmup. 

Hi @ahatamiz ,

Thanks for providing the nice `UNETR` feature in MONAI v0.6.
Could you please help self-assign this ticket?

Thanks in advance.

Hi @Nic-Ma 
Sure, however I don't think I have been given permission to self-assign. I don't see the option yet. Would you mind provide that permission ? 

I will submit a new PR regarding this. 

Thanks,
Ali 
Hi @ahatamiz ,

I assigned the ticket to you.

Thanks in advance.

Thank you @Nic-Ma 

Resolving patch:
diff --git a/monai/networks/nets/unetr.py b/monai/networks/nets/unetr.py
--- a/monai/networks/nets/unetr.py
+++ b/monai/networks/nets/unetr.py
@@ -29,14 +29,14 @@ def __init__(
         in_channels: int,
         out_channels: int,
         img_size: Tuple[int, int, int],
-        feature_size: int,
-        hidden_size: int,
-        mlp_dim: int,
-        num_heads: int,
-        pos_embed: str,
-        norm_name: Union[Tuple, str],
+        feature_size: int = 16,
+        hidden_size: int = 768,
+        mlp_dim: int = 3072,
+        num_heads: int = 12,
+        pos_embed: str = "perceptron",
+        norm_name: Union[Tuple, str] = "instance",
         conv_block: bool = False,
-        res_block: bool = False,
+        res_block: bool = True,
         dropout_rate: float = 0.0,
     ) -> None:
         """
@@ -54,6 +54,14 @@ def __init__(
             res_block: bool argument to determine if residual block is used.
             dropout_rate: faction of the input units to drop.
 
+        Examples::
+
+            # for single channel input 4-channel output with patch 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 4-channel input 3-channel output with patch 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')
+
         """
 
         super().__init__()
diff --git a/monai/networks/nets/vit.py b/monai/networks/nets/vit.py
--- a/monai/networks/nets/vit.py
+++ b/monai/networks/nets/vit.py
@@ -29,12 +29,12 @@ def __init__(
         in_channels: int,
         img_size: Tuple[int, int, int],
         patch_size: Tuple[int, int, int],
-        hidden_size: int,
-        mlp_dim: int,
-        num_layers: int,
-        num_heads: int,
-        pos_embed: str,
-        classification: bool,
+        hidden_size: int = 768,
+        mlp_dim: int = 3072,
+        num_layers: int = 12,
+        num_heads: int = 12,
+        pos_embed: str = "perceptron",
+        classification: bool = False,
         num_classes: int = 2,
         dropout_rate: float = 0.0,
     ) -> None:
@@ -52,6 +52,14 @@ def __init__(
             num_classes: number of classes if classification is used.
             dropout_rate: faction of the input units to drop.
 
+        Examples::
+
+            # for single channel input with patch 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
+            >>> net = ViT(in_channels=3, img_size=(128,128,128), pos_embed='conv', classification= True)
+
         """
 
         super().__init__()
diff --git a/monai/optimizers/lr_scheduler.py b/monai/optimizers/lr_scheduler.py
--- a/monai/optimizers/lr_scheduler.py
+++ b/monai/optimizers/lr_scheduler.py
@@ -9,8 +9,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import math
+
 from torch.optim import Optimizer
-from torch.optim.lr_scheduler import _LRScheduler
+from torch.optim.lr_scheduler import LambdaLR, _LRScheduler
 
 __all__ = ["LinearLR", "ExponentialLR"]
 
@@ -52,3 +54,33 @@ class ExponentialLR(_LRSchedulerMONAI):
     def get_lr(self):
         r = self.last_epoch / (self.num_iter - 1)
         return [base_lr * (self.end_lr / base_lr) ** r for base_lr in self.base_lrs]
+
+
+class WarmupCosineSchedule(LambdaLR):
+    """Linear warmup and then cosine decay.
+    Based on https://huggingface.co/ implementation.
+    """
+
+    def __init__(
+        self, optimizer: Optimizer, warmup_steps: int, t_total: int, cycles: float = 0.5, last_epoch: int = -1
+    ) -> None:
+        """
+        Args:
+            optimizer: wrapped optimizer.
+            warmup_steps: number of warmup iterations.
+            t_total: total number of training iterations.
+            cycles: cosine cycles parameter.
+            last_epoch: the index of last epoch.
+        Returns:
+            None
+        """
+        self.warmup_steps = warmup_steps
+        self.t_total = t_total
+        self.cycles = cycles
+        super(WarmupCosineSchedule, self).__init__(optimizer, self.lr_lambda, last_epoch)
+
+    def lr_lambda(self, step):
+        if step < self.warmup_steps:
+            return float(step) / float(max(1.0, self.warmup_steps))
+        progress = float(step - self.warmup_steps) / float(max(1, self.t_total - self.warmup_steps))
+        return max(0.0, 0.5 * (1.0 + math.cos(math.pi * float(self.cycles) * 2.0 * progress)))


Test patch:
diff --git a/tests/test_lr_scheduler.py b/tests/test_lr_scheduler.py
new file mode 100644
--- /dev/null
+++ b/tests/test_lr_scheduler.py
@@ -0,0 +1,58 @@
+# Copyright 2020 - 2021 MONAI Consortium
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#     http://www.apache.org/licenses/LICENSE-2.0
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+
+import torch
+from parameterized import parameterized
+
+from monai.optimizers.lr_scheduler import WarmupCosineSchedule
+
+
+class SchedulerTestNet(torch.nn.Module):
+    def __init__(self):
+        super(SchedulerTestNet, self).__init__()
+        self.conv1 = torch.nn.Conv2d(1, 1, 1)
+        self.conv2 = torch.nn.Conv2d(1, 1, 1)
+
+    def forward(self, x):
+        return self.conv2(torch.nn.functional.relu(self.conv1(x)))
+
+
+TEST_CASE_LRSCHEDULER = [
+    [
+        {
+            "warmup_steps": 2,
+            "t_total": 10,
+        },
+        [0.000, 0.500, 1.00, 0.962, 0.854, 0.691, 0.500, 0.309, 0.146, 0.038],
+    ]
+]
+
+
+class TestLRSCHEDULER(unittest.TestCase):
+    @parameterized.expand(TEST_CASE_LRSCHEDULER)
+    def test_shape(self, input_param, expected_lr):
+        net = SchedulerTestNet()
+        optimizer = torch.optim.Adam(net.parameters(), lr=1.0)
+        scheduler = WarmupCosineSchedule(optimizer, **input_param)
+        self.assertEqual(len([scheduler.get_last_lr()[0]]), 1)
+        lrs_1 = []
+        for _ in range(input_param["t_total"]):
+            lrs_1.append(float("{:.3f}".format(scheduler.get_last_lr()[0])))
+            optimizer.step()
+            scheduler.step()
+        for a, b in zip(lrs_1, expected_lr):
+            self.assertEqual(a, b, msg="LR is wrong ! expected {}, got {}".format(b, a))
+
+
+if __name__ == "__main__":
+    unittest.main()