feat: Added the contiguous param tricks to training scripts

references/classification/train.py

@@ -20,6 +20,7 @@
 from torch import nn
 import torchvision
 from torchvision import transforms
+from contiguous_params import ContiguousParams
 
 import holocron
 
@@ -157,17 +158,19 @@ def main(args):
     elif args.loss == 'label_smoothing':
         criterion = holocron.nn.LabelSmoothingCrossEntropy()
 
+    # Create the contiguous parameters.
+    model_params = ContiguousParams([p for p in model.parameters() if p.requires_grad])
     if args.opt == 'adam':
-        optimizer = torch.optim.Adam([p for p in model.parameters() if p.requires_grad], args.lr,
+        optimizer = torch.optim.Adam(model_params.contiguous(), args.lr,
                                      betas=(0.95, 0.99), eps=1e-6, weight_decay=args.weight_decay)
     elif args.opt == 'radam':
-        optimizer = holocron.optim.RAdam([p for p in model.parameters() if p.requires_grad], args.lr,
+        optimizer = holocron.optim.RAdam(model_params.contiguous(), args.lr,
                                          betas=(0.95, 0.99), eps=1e-6, weight_decay=args.weight_decay)
     elif args.opt == 'ranger':
-        optimizer = Lookahead(holocron.optim.RAdam([p for p in model.parameters() if p.requires_grad], args.lr,
+        optimizer = Lookahead(holocron.optim.RAdam(model_params.contiguous(), args.lr,
                                                    betas=(0.95, 0.99), eps=1e-6, weight_decay=args.weight_decay))
     elif args.opt == 'tadam':
-        optimizer = holocron.optim.TAdam([p for p in model.parameters() if p.requires_grad], args.lr,
+        optimizer = holocron.optim.TAdam(model_params.contiguous(), args.lr,
                                          betas=(0.95, 0.99), eps=1e-6, weight_decay=args.weight_decay)
 
     if args.lr_finder:
@@ -200,6 +203,8 @@ def main(args):
     mb = master_bar(range(args.start_epoch, args.epochs))
     for epoch in mb:
         train_one_epoch(model, criterion, optimizer, lr_scheduler, train_loader, device, mb)
+        # Check that the optimizer only applies valid ops.
+        model_params.assert_buffer_is_valid()
         val_loss, acc1, acc5 = evaluate(model, criterion, val_loader, device=device)
         if args.sched == 'plateau':
             lr_scheduler.step(val_loss)

references/detection/train.py

@@ -19,6 +19,7 @@
 from torchvision.datasets import VOCDetection
 from torchvision.ops.boxes import box_iou
 from torchvision.transforms import functional as F
+from contiguous_params import ContiguousParams
 
 import holocron
 from transforms import (Compose, VOCTargetTransform, Resize, ImageTransform, CenterCrop, RandomResizedCrop,
@@ -231,17 +232,18 @@ def main(args):
             p.requires_grad_(False)
     model.to(device)
 
+    model_params = ContiguousParams([p for p in model.parameters() if p.requires_grad])
     if args.opt == 'adam':
-        optimizer = torch.optim.Adam([p for p in model.parameters() if p.requires_grad], args.lr,
+        optimizer = torch.optim.Adam(model_params.contiguous(), args.lr,
                                      betas=(0.95, 0.99), eps=1e-6, weight_decay=args.weight_decay)
     elif args.opt == 'radam':
-        optimizer = holocron.optim.RAdam([p for p in model.parameters() if p.requires_grad], args.lr,
+        optimizer = holocron.optim.RAdam(model_params.contiguous(), args.lr,
                                          betas=(0.95, 0.99), eps=1e-6, weight_decay=args.weight_decay)
     elif args.opt == 'ranger':
-        optimizer = Lookahead(holocron.optim.RAdam([p for p in model.parameters() if p.requires_grad], args.lr,
+        optimizer = Lookahead(holocron.optim.RAdam(model_params.contiguous(), args.lr,
                                                    betas=(0.95, 0.99), eps=1e-6, weight_decay=args.weight_decay))
     elif args.opt == 'tadam':
-        optimizer = holocron.optim.TAdam([p for p in model.parameters() if p.requires_grad], args.lr,
+        optimizer = holocron.optim.TAdam(model_params.contiguous(), args.lr,
                                          betas=(0.95, 0.99), eps=1e-6, weight_decay=args.weight_decay)
 
     if args.lr_finder:
@@ -276,6 +278,8 @@ def main(args):
     mb = master_bar(range(args.start_epoch, args.epochs))
     for epoch in mb:
         train_one_epoch(model, optimizer, lr_scheduler, train_loader, device, mb)
+        # Check that the optimizer only applies valid ops.
+        model_params.assert_buffer_is_valid()
         loc_err, clf_err, det_err = evaluate(model, val_loader, device=device)
         mb.main_bar.comment = f"Epoch {args.start_epoch+epoch+1}/{args.start_epoch+args.epochs}"
         mb.write(f"Epoch {args.start_epoch+epoch+1}/{args.start_epoch+args.epochs} - "

references/requirements.txt

@@ -0,0 +1,5 @@
+pylocron
+tqdm
+fastprogress>=0.1.21
+matplotlib>=3.0.0
+contiguous_params @ git+https://github.com/philjd/contiguous_pytorch_params.git#egg=contiguous_params

references/segmentation/train.py

@@ -18,6 +18,7 @@
 from torchvision.datasets import VOCSegmentation
 from torchvision.ops.misc import FrozenBatchNorm2d
 from torchvision.transforms import functional as F
+from contiguous_params import ContiguousParams
 
 import holocron
 from transforms import (Compose, Resize, ImageTransform, CenterCrop, RandomResizedCrop,
@@ -192,14 +193,15 @@ def main(args):
             p.requires_grad_(False)
     model.to(device)
 
+    model_params = ContiguousParams([p for p in model.parameters() if p.requires_grad])
     if args.opt == 'adam':
-        optimizer = torch.optim.Adam(model.parameters(), args.lr, betas=(0.95, 0.99), eps=1e-6,
+        optimizer = torch.optim.Adam(model_params.contiguous(), args.lr, betas=(0.95, 0.99), eps=1e-6,
                                      weight_decay=args.weight_decay)
     elif args.opt == 'radam':
-        optimizer = holocron.optim.RAdam(model.parameters(), args.lr, betas=(0.95, 0.99), eps=1e-6,
+        optimizer = holocron.optim.RAdam(model_params.contiguous(), args.lr, betas=(0.95, 0.99), eps=1e-6,
                                          weight_decay=args.weight_decay)
     elif args.opt == 'ranger':
-        optimizer = Lookahead(holocron.optim.RAdam(model.parameters(), args.lr, betas=(0.95, 0.99), eps=1e-6,
+        optimizer = Lookahead(holocron.optim.RAdam(model_params.contiguous(), args.lr, betas=(0.95, 0.99), eps=1e-6,
                                                    weight_decay=args.weight_decay))
 
     loss_weight = torch.ones(len(classes))
@@ -238,6 +240,8 @@ def main(args):
     mb = master_bar(range(args.start_epoch, args.epochs))
     for epoch in mb:
         train_one_epoch(model, optimizer, criterion, lr_scheduler, train_loader, device, mb)
+        # Check that the optimizer only applies valid ops.
+        model_params.assert_buffer_is_valid()
         val_loss, mean_iou = evaluate(model, val_loader, criterion, device=device)
         mb.main_bar.comment = f"Epoch {args.start_epoch+epoch+1}/{args.start_epoch+args.epochs}"
         mb.write(f"Epoch {args.start_epoch+epoch+1}/{args.start_epoch+args.epochs} - "