Added SWA to test

Author: vfdev-5

README.md

@@ -77,7 +77,7 @@ mlflow run experiments/ --experiment-name=CIFAR10 -P dataset=CIFAR10 -P network=
 ```
 export MLFLOW_TRACKING_URI=$OUTPUT_PATH/mlruns
 
-mlflow run experiments/ --experiment-name=CIFAR10 -P dataset=CIFAR10 -P network=wideresnet -P params="data_path=../input/cifar10;num_epochs=6250;learning_rate=0.03;batch_size=64;TSA_proba_min=0.1;unlabelled_batch_size=320;"
+mlflow run experiments/ --experiment-name=CIFAR10 -P dataset=CIFAR10 -P network=wideresnet -P params="data_path=../input/cifar10;num_epochs=6250;learning_rate=0.03;batch_size=64;TSA_proba_min=0.1;unlabelled_batch_size=320;num_warmup_steps=20000"
 ```
 
 Unfortunately, I can not reproduce paper's result with 5.3 test error.

code/main.py

@@ -16,6 +16,8 @@
 import torch.optim as optim
 from torch.optim.lr_scheduler import CosineAnnealingLR
 
+import torchcontrib
+
 import ignite
 from ignite.engine import Events, Engine, create_supervised_evaluator
 from ignite.metrics import Accuracy, Loss, RunningAverage
@@ -53,7 +55,11 @@ def run(output_path, config):
                           momentum=config['momentum'],
                           weight_decay=config['weight_decay'],
                           nesterov=True)
-    
+
+    with_SWA = config['with_SWA']
+    if with_SWA:
+        optimizer = torchcontrib.optim.SWA(optimizer)
+
     criterion = nn.CrossEntropyLoss().to(device)
     if config['consistency_criterion'] == "MSE":
         consistency_criterion = nn.MSELoss()
@@ -177,6 +183,17 @@ def log_learning_rate(engine):
             lr = optimizer.param_groups[0]['lr']
             mlflow.log_metric("learning rate", lr, step=step)
 
+    if with_SWA:
+        @trainer.on(Events.COMPLETED)
+        def swap_swa_sgd(engine):
+            optimizer.swap_swa_sgd()
+            optimizer.bn_update(train_labelled_loader, model)
+
+        @trainer.on(Events.EPOCH_COMPLETED)
+        def update_swa(engine):
+            if engine.state.epoch - 1 > int(num_epochs * 0.75): 
+                optimizer.update_swa()
+
     metric_names = [
         'supervised batch loss',
         'consistency batch loss',
@@ -301,6 +318,8 @@ def mlflow_val_metrics_logging(engine, tag):
         "TSA_proba_max": 1.0,
 
         "no_UDA": False,  # disable UDA training 
+
+        "with_SWA": False,
     }
 
     # Override config:

experiments/conda.yaml

@@ -14,3 +14,4 @@ dependencies:
     - tensorboardX
     - nested_dict
     - git+https://github.com/pytorch/ignite.git
+    - torchcontrib