0. Experimental Setup

0-1. Prepare Library

pipinstall-rrequirements.txt

0-2. Download dataset (Kaggle Intel Image Classification)

1. Baseline Training

pythonmain.py--checkpoint_namebaseline;

2. Stochastic Weight Averaging Training

fromtorch.optim.swa_utilsimportAveragedModel,SWALRfromtorch.optim.lr_schedulerimportCosineAnnealingLRloader,optimizer,model,loss_fn= ...swa_model=AveragedModel(model)scheduler=CosineAnnealingLR(optimizer,T_max=100)swa_start=5swa_scheduler=SWALR(optimizer,swa_lr=0.05)forepochinrange(100):forinput,targetinloader:optimizer.zero_grad()loss_fn(model(input),target).backward()optimizer.step()ifepoch>swa_start:swa_model.update_parameters(model)swa_scheduler.step()else:scheduler.step()# Update bn statistics for the swa_model at the endtorch.optim.swa_utils.update_bn(loader,swa_model)# Use swa_model to make predictions on test datapreds=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.pyforbatch_idx, (inputs,labels)inenumerate(data_loader):ifnotargs.decay_type=='swa':self.scheduler.step()else:ifepoch<=args.swa_start:self.scheduler.step()ifepoch>args.swa_startandargs.decay_type=='swa':self.swa_model.update_parameters(self.model)self.swa_scheduler.step()...# in main.pyswa_model=swa_model.cpu()torch.optim.swa_utils.update_bn(train_loader,swa_model)swa_model=swa_model.cuda()

Run Script (Command Line)

pythonmain.py--checkpoint_nameswa--decay_typeswa--swa_start90--swa_lr5e-5;

3. Performance Table

4. Code Reference