Add crop_mode for pretraind config / image transforms. Add support for dynamo compilation to benchmark/train/validate

Author: rwightman

benchmark.py

@@ -56,6 +56,13 @@
 except ImportError as e:
     has_functorch = False
 
+try:
+    import torch._dynamo
+    has_dynamo = True
+except ImportError:
+    has_dynamo = False
+    pass
+
 
 if torch.cuda.is_available():
     torch.backends.cuda.matmul.allow_tf32 = True
@@ -106,13 +113,19 @@
                     help='Numeric precision. One of (amp, float32, float16, bfloat16, tf32)')
 parser.add_argument('--fuser', default='', type=str,
                     help="Select jit fuser. One of ('', 'te', 'old', 'nvfuser')")
+parser.add_argument('--dynamo-backend', default=None, type=str,
+                    help="Select dynamo backend. Default: None")
+parser.add_argument('--fast-norm', default=False, action='store_true',
+                    help='enable experimental fast-norm')
+
+# codegen (model compilation) options
 scripting_group = parser.add_mutually_exclusive_group()
 scripting_group.add_argument('--torchscript', dest='torchscript', action='store_true',
-                    help='convert model torchscript for inference')
+                             help='convert model torchscript for inference')
 scripting_group.add_argument('--aot-autograd', default=False, action='store_true',
-                    help="Enable AOT Autograd support. (It's recommended to use this option with `--fuser nvfuser` together)")
-scripting_group.add_argument('--fast-norm', default=False, action='store_true',
-                    help='enable experimental fast-norm')
+                             help="Enable AOT Autograd optimization.")
+scripting_group.add_argument('--dynamo', default=False, action='store_true',
+                             help="Enable Dynamo optimization.")
 
 # train optimizer parameters
 parser.add_argument('--opt', default='sgd', type=str, metavar='OPTIMIZER',
@@ -206,6 +219,8 @@ def __init__(
             device='cuda',
             torchscript=False,
             aot_autograd=False,
+            dynamo=False,
+            dynamo_backend=None,
             precision='float32',
             fuser='',
             num_warm_iter=10,
@@ -241,14 +256,21 @@ def __init__(
         _logger.info('Model %s created, param count: %d' % (model_name, self.param_count))
 
         data_config = resolve_data_config(kwargs, model=self.model, use_test_size=not use_train_size)
+        self.input_size = data_config['input_size']
+        self.batch_size = kwargs.pop('batch_size', 256)
+
         self.scripted = False
         if torchscript:
             self.model = torch.jit.script(self.model)
             self.scripted = True
-        self.input_size = data_config['input_size']
-        self.batch_size = kwargs.pop('batch_size', 256)
-
-        if aot_autograd:
+        elif dynamo:
+            assert has_dynamo, "torch._dynamo is needed for --dynamo"
+            torch._dynamo.reset()
+            if dynamo_backend is not None:
+                self.model = torch._dynamo.optimize(dynamo_backend)(self.model)
+            else:
+                self.model = torch._dynamo.optimize()(self.model)
+        elif aot_autograd:
             assert has_functorch, "functorch is needed for --aot-autograd"
             self.model = memory_efficient_fusion(self.model)
 

timm/data/config.py

@@ -5,9 +5,15 @@
 _logger = logging.getLogger(__name__)
 
 
-def resolve_data_config(args, default_cfg={}, model=None, use_test_size=False, verbose=False):
+def resolve_data_config(
+        args,
+        default_cfg=None,
+        model=None,
+        use_test_size=False,
+        verbose=False
+):
     new_config = {}
-    default_cfg = default_cfg
+    default_cfg = default_cfg or {}
     if not default_cfg and model is not None and hasattr(model, 'default_cfg'):
         default_cfg = model.default_cfg
 
@@ -63,7 +69,7 @@ def resolve_data_config(args, default_cfg={}, model=None, use_test_size=False, v
     elif default_cfg.get('std', None):
         new_config['std'] = default_cfg['std']
 
-    # resolve default crop percentage
+    # resolve default inference crop
     crop_pct = DEFAULT_CROP_PCT
     if args.get('crop_pct', None):
         crop_pct = args['crop_pct']
@@ -74,6 +80,14 @@ def resolve_data_config(args, default_cfg={}, model=None, use_test_size=False, v
             crop_pct = default_cfg['crop_pct']
     new_config['crop_pct'] = crop_pct
 
+    # resolve default crop percentage
+    crop_mode = DEFAULT_CROP_MODE
+    if args.get('crop_mode', None):
+        crop_mode = args['crop_mode']
+    elif default_cfg.get('crop_mode', None):
+        crop_mode = default_cfg['crop_mode']
+    new_config['crop_mode'] = crop_mode
+
     if verbose:
         _logger.info('Data processing configuration for current model + dataset:')
         for n, v in new_config.items():

timm/data/constants.py

@@ -1,4 +1,5 @@
 DEFAULT_CROP_PCT = 0.875
+DEFAULT_CROP_MODE = 'center'
 IMAGENET_DEFAULT_MEAN = (0.485, 0.456, 0.406)
 IMAGENET_DEFAULT_STD = (0.229, 0.224, 0.225)
 IMAGENET_INCEPTION_MEAN = (0.5, 0.5, 0.5)

timm/data/loader.py

@@ -211,6 +211,7 @@ def create_loader(
         num_workers=1,
         distributed=False,
         crop_pct=None,
+        crop_mode=None,
         collate_fn=None,
         pin_memory=False,
         fp16=False,  # deprecated, use img_dtype
@@ -240,6 +241,7 @@ def create_loader(
         mean=mean,
         std=std,
         crop_pct=crop_pct,
+        crop_mode=crop_mode,
         tf_preprocessing=tf_preprocessing,
         re_prob=re_prob,
         re_mode=re_mode,

timm/data/tf_preprocessing.py

@@ -22,12 +22,13 @@
 # limitations under the License.
 # ==============================================================================
 """ImageNet preprocessing for MnasNet."""
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 import numpy as np
 
 IMAGE_SIZE = 224
 CROP_PADDING = 32
 
+tf.compat.v1.disable_eager_execution()
 
 def distorted_bounding_box_crop(image_bytes,
                                 bbox,

timm/data/transforms.py

@@ -1,3 +1,9 @@
+import math
+import numbers
+import random
+import warnings
+from typing import List, Sequence
+
 import torch
 import torchvision.transforms.functional as F
 try:
@@ -6,9 +12,6 @@
 except ImportError:
     has_interpolation_mode = False
 from PIL import Image
-import warnings
-import math
-import random
 import numpy as np
 
 
@@ -96,6 +99,19 @@ def interp_mode_to_str(mode):
 _RANDOM_INTERPOLATION = (str_to_interp_mode('bilinear'), str_to_interp_mode('bicubic'))
 
 
+def _setup_size(size, error_msg):
+    if isinstance(size, numbers.Number):
+        return int(size), int(size)
+
+    if isinstance(size, Sequence) and len(size) == 1:
+        return size[0], size[0]
+
+    if len(size) != 2:
+        raise ValueError(error_msg)
+
+    return size
+
+
 class RandomResizedCropAndInterpolation:
     """Crop the given PIL Image to random size and aspect ratio with random interpolation.
 
@@ -195,3 +211,132 @@ def __repr__(self):
         format_string += ', ratio={0}'.format(tuple(round(r, 4) for r in self.ratio))
         format_string += ', interpolation={0})'.format(interpolate_str)
         return format_string
+
+
+def center_crop_or_pad(img: torch.Tensor, output_size: List[int], fill=0) -> torch.Tensor:
+    """Center crops and/or pads the given image.
+    If the image is torch Tensor, it is expected
+    to have [..., H, W] shape, where ... means an arbitrary number of leading dimensions.
+    If image size is smaller than output size along any edge, image is padded with 0 and then center cropped.
+
+    Args:
+        img (PIL Image or Tensor): Image to be cropped.
+        output_size (sequence or int): (height, width) of the crop box. If int or sequence with single int,
+            it is used for both directions.
+        fill (int, Tuple[int]): Padding color
+
+    Returns:
+        PIL Image or Tensor: Cropped image.
+    """
+    if isinstance(output_size, numbers.Number):
+        output_size = (int(output_size), int(output_size))
+    elif isinstance(output_size, (tuple, list)) and len(output_size) == 1:
+        output_size = (output_size[0], output_size[0])
+
+    _, image_height, image_width = F.get_dimensions(img)
+    crop_height, crop_width = output_size
+
+    if crop_width > image_width or crop_height > image_height:
+        padding_ltrb = [
+            (crop_width - image_width) // 2 if crop_width > image_width else 0,
+            (crop_height - image_height) // 2 if crop_height > image_height else 0,
+            (crop_width - image_width + 1) // 2 if crop_width > image_width else 0,
+            (crop_height - image_height + 1) // 2 if crop_height > image_height else 0,
+        ]
+        img = F.pad(img, padding_ltrb, fill=fill)
+        _, image_height, image_width = F.get_dimensions(img)
+        if crop_width == image_width and crop_height == image_height:
+            return img
+
+    crop_top = int(round((image_height - crop_height) / 2.0))
+    crop_left = int(round((image_width - crop_width) / 2.0))
+    return F.crop(img, crop_top, crop_left, crop_height, crop_width)
+
+
+class CenterCropOrPad(torch.nn.Module):
+    """Crops the given image at the center.
+    If the image is torch Tensor, it is expected
+    to have [..., H, W] shape, where ... means an arbitrary number of leading dimensions.
+    If image size is smaller than output size along any edge, image is padded with 0 and then center cropped.
+
+    Args:
+        size (sequence or int): Desired output size of the crop. If size is an
+            int instead of sequence like (h, w), a square crop (size, size) is
+            made. If provided a sequence of length 1, it will be interpreted as (size[0], size[0]).
+    """
+
+    def __init__(self, size, fill=0):
+        super().__init__()
+        self.size = _setup_size(size, error_msg="Please provide only two dimensions (h, w) for size.")
+        self.fill = fill
+
+    def forward(self, img):
+        """
+        Args:
+            img (PIL Image or Tensor): Image to be cropped.
+
+        Returns:
+            PIL Image or Tensor: Cropped image.
+        """
+        return center_crop_or_pad(img, self.size, fill=self.fill)
+
+    def __repr__(self) -> str:
+        return f"{self.__class__.__name__}(size={self.size})"
+
+
+class ResizeKeepRatio:
+    """ Resize and Keep Ratio
+    """
+
+    def __init__(
+            self,
+            size,
+            longest=0.,
+            interpolation='bilinear',
+            fill=0,
+    ):
+        if isinstance(size, (list, tuple)):
+            self.size = tuple(size)
+        else:
+            self.size = (size, size)
+        self.interpolation = str_to_interp_mode(interpolation)
+        self.longest = float(longest)
+        self.fill = fill
+
+    @staticmethod
+    def get_params(img, target_size, longest):
+        """Get parameters
+
+        Args:
+            img (PIL Image): Image to be cropped.
+            target_size (Tuple[int, int]): Size of output
+        Returns:
+            tuple: params (h, w) and (l, r, t, b) to be passed to ``resize`` and ``pad`` respectively
+        """
+        source_size = img.size[::-1]  # h, w
+        h, w = source_size
+        target_h, target_w = target_size
+        ratio_h = h / target_h
+        ratio_w = w / target_w
+        ratio = max(ratio_h, ratio_w) * longest + min(ratio_h, ratio_w) * (1. - longest)
+        size = [round(x / ratio) for x in source_size]
+        return size
+
+    def __call__(self, img):
+        """
+        Args:
+            img (PIL Image): Image to be cropped and resized.
+
+        Returns:
+            PIL Image: Resized, padded to at least target size, possibly cropped to exactly target size
+        """
+        size = self.get_params(img, self.size, self.longest)
+        img = F.resize(img, size, self.interpolation)
+        return img
+
+    def __repr__(self):
+        interpolate_str = interp_mode_to_str(self.interpolation)
+        format_string = self.__class__.__name__ + '(size={0}'.format(self.size)
+        format_string += f', interpolation={interpolate_str})'
+        format_string += f', longest={self.longest:.3f})'
+        return format_string