Improve vision classification (#268)

* Added PiecewiseLinear Scheduler

* Fix typo

* Modified the Transformations

* Modifies so that lint is satisfactory

* Fix typo

Author: theory-in-progress

src/templates/template-vision-classification/data.py

@@ -15,10 +15,20 @@ def setup_data(config: Any):
     #::: if (it.use_dist) { :::#
     local_rank = idist.get_local_rank()
     #::: } :::#
-    transform = T.Compose(
+    train_transform = T.Compose(
         [
+            T.Pad(4),
+            T.RandomCrop(32, fill=128),
+            T.RandomHorizontalFlip(),
             T.ToTensor(),
-            T.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
+            T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+        ]
+    )
+
+    eval_transform = T.Compose(
+        [
+            T.ToTensor(),
+            T.Normalize(mean=[0.4914, 0.4822, 0.4465], std=[0.2023, 0.1994, 0.2010]),
         ]
     )
 
@@ -32,13 +42,13 @@ def setup_data(config: Any):
         root=config.data_path,
         train=True,
         download=True,
-        transform=transform,
+        transform=train_transform,
     )
     dataset_eval = torchvision.datasets.CIFAR10(
         root=config.data_path,
         train=False,
         download=True,
-        transform=transform,
+        transform=eval_transform,
     )
 
     #::: if (it.use_dist) { :::#

src/templates/template-vision-classification/main.py

@@ -5,6 +5,7 @@
 import yaml
 from data import setup_data
 from ignite.engine import Events
+from ignite.handlers import PiecewiseLinear
 from ignite.metrics import Accuracy, Loss
 from ignite.utils import manual_seed
 from models import setup_model
@@ -29,6 +30,15 @@ def run(local_rank: int, config: Any):
     model = idist.auto_model(setup_model(config.model))
     optimizer = idist.auto_optim(optim.Adam(model.parameters(), lr=config.lr))
     loss_fn = nn.CrossEntropyLoss().to(device=device)
+    milestones_values = [
+        (0, 0.0),
+        (
+            len(dataloader_train),
+            config.lr,
+        ),
+        (config.max_epochs * len(dataloader_train), 0.0),
+    ]
+    lr_scheduler = PiecewiseLinear(optimizer, "lr", milestones_values=milestones_values)
 
     # trainer and evaluator
     trainer = setup_trainer(
@@ -53,10 +63,17 @@ def run(local_rank: int, config: Any):
     (config.output_dir / "config-lock.yaml").write_text(yaml.dump(config))
     trainer.logger = evaluator.logger = logger
 
+    trainer.add_event_handler(Events.ITERATION_COMPLETED, lr_scheduler)
+
     # setup ignite handlers
     #::: if (it.save_training || it.save_evaluation) { :::#
     #::: if (it.save_training) { :::#
-    to_save_train = {"model": model, "optimizer": optimizer, "trainer": trainer}
+    to_save_train = {
+        "model": model,
+        "optimizer": optimizer,
+        "trainer": trainer,
+        "lr_scheduler": lr_scheduler,
+    }
     #::: } else { :::#
     to_save_train = None
     #::: } :::#