Add DynamicResolution class and enable dynamic resolution option

vitookit/datasets/ffcv_transform.py

@@ -35,9 +35,49 @@ from torch import nn
 IMAGENET_MEAN = np.array([0.485, 0.456, 0.406]) * 255
 IMAGENET_STD = np.array([0.229, 0.224, 0.225]) * 255
 
+@gin.configurable
+class DynamicResolution:
+    def __init__(self, start_ramp=gin.REQUIRED, end_ramp=gin.REQUIRED,  min_res=160, max_res=224, step=32):        
+        self.start_ramp = start_ramp
+        self.end_ramp = end_ramp
+        self.min_res = min_res
+        self.max_res = max_res
+        self.step = step
+        
+        assert min_res <= max_res
+        
+        
+    def get_resolution(self, epoch):    
+        if epoch <= self.start_ramp:
+            return self.min_res
 
+        if epoch >= self.end_ramp:
+            return self.max_res
+
+        # otherwise, linearly interpolate to the nearest multiple of 32
+        interp = np.interp([epoch], [self.start_ramp, self.end_ramp], [self.min_res, self.max_res])
+        final_res = int(np.round(interp[0] / self.step)) * self.step
+        return final_res
+    
+    def __call__(self, loader, epoch,is_ffcv=False):
+        img_size = self.get_resolution(epoch)
+        print("resolution = %s" % str(img_size))
+        if is_ffcv:
+            pipeline=loader.pipeline_specs['image']
+            if pipeline.decoder.output_size[0] != img_size:
+                k = img_size/self.max_res
+                pipeline.decoder.scale = (0.2*k,1)
+                pipeline.decoder.output_size = (img_size,img_size)
+                loader.generate_code()
+        else:
+            decoder = loader.dataset.transforms.transform.transforms[0]
+            k = img_size/self.max_res
+            decoder.size=(img_size,img_size)
+            decoder.scale = (0.2*k,1)
+                
+                
 @gin.configurable
-def SimplePipeline(img_size=224,scale=(0.2,1), ratio=(3.0/4.0, 4.0/3.0),blur=False):
+def SimplePipeline(img_size=224,scale=(0.2,1), ratio=(3.0/4.0, 4.0/3.0)):
     image_pipeline = [
             RandomResizedCropRGBImageDecoder((img_size, img_size), scale=scale,ratio=ratio),
             RandomHorizontalFlip(),
@@ -46,8 +86,6 @@ def SimplePipeline(img_size=224,scale=(0.2,1), ratio=(3.0/4.0, 4.0/3.0),blur=Fal
             ToDevice(torch.device('cuda')),        
             ToTorchImage(),
             ]
-    if blur:
-        image_pipeline.append(transforms.GaussianBlur(3))
     label_pipeline = [IntDecoder(), ToTensor(),ToDevice(torch.device('cuda'))]
     # Pipeline for each data field
     pipelines = {

vitookit/evaluation/eval_cls_ffcv.py

@@ -51,6 +51,7 @@ def get_args_parser():
                         help='Accumulate gradient iterations (for increasing the effective batch size under memory constraints)')
     parser.add_argument('--epochs', default=100, type=int)
     parser.add_argument('--ckpt_freq', default=5, type=int)
+    parser.add_argument("--dynamic_resolution", default=False, action="store_true", help="Use dynamic resolution.")
 
     # Model parameters
     parser.add_argument("--compile", action='store_true', default=False, help="compile model with PyTorch 2.0")
@@ -354,6 +355,14 @@ def main(args):
         exit(0)
 
     print(f"Start training for {args.epochs} epochs from {args.start_epoch}")
+    
+    if args.dynamic_resolution:
+        import torch._dynamo
+        torch._dynamo.config.suppress_errors = True
+        dres = DynamicResolution()    
+    else:
+        dres = None
+        
     start_time = time.time()
     max_accuracy = 0.0
     if args.output_dir and misc.is_main_process():
@@ -364,7 +373,8 @@ def main(args):
             pass
         
     for epoch in range(args.start_epoch, args.epochs):
-
+        if dres:
+            dres(data_loader_train,epoch,True)
         train_stats = train_one_epoch(
             model, criterion, data_loader_train,
             optimizer, device, epoch, loss_scaler,lr_scheduler,