pip install -r requirements.txt

python main.py --checkpoint_name baseline;

from torch.optim.swa_utils import AveragedModel, SWALR
from torch.optim.lr_scheduler import CosineAnnealingLR

loader, optimizer, model, loss_fn = ...
swa_model = AveragedModel(model)
scheduler = CosineAnnealingLR(optimizer, T_max=100)
swa_start = 5
swa_scheduler = SWALR(optimizer, swa_lr=0.05)

for epoch in range(100):
      for input, target in loader:
          optimizer.zero_grad()
          loss_fn(model(input), target).backward()
          optimizer.step()
      if epoch > swa_start:
          swa_model.update_parameters(model)
          swa_scheduler.step()
      else:
          scheduler.step()

# Update bn statistics for the swa_model at the end
torch.optim.swa_utils.update_bn(loader, swa_model)
# Use swa_model to make predictions on test data 
preds = swa_model(test_input)

# in main.py
""" define model and learning rate scheduler for stochastic weight averaging """
swa_model = torch.optim.swa_utils.AveragedModel(model)
swa_scheduler = SWALR(optimizer, swa_lr=args.swa_lr)

...  

# in learning/trainer.py
for batch_idx, (inputs, labels) in enumerate(data_loader):
  if not args.decay_type == 'swa':
        self.scheduler.step()
  else:
      if epoch <= args.swa_start:
          self.scheduler.step()

if epoch > args.swa_start and args.decay_type == 'swa':
  self.swa_model.update_parameters(self.model)
  self.swa_scheduler.step()

...

# in main.py
swa_model = swa_model.cpu()
torch.optim.swa_utils.update_bn(train_loader, swa_model)
swa_model = swa_model.cuda() 

python main.py --checkpoint_name swa --decay_type swa --swa_start 90 --swa_lr 5e-5;