Allow setting image size in reference training.

Author: hgaiser

references/detection/presets.py

@@ -15,9 +15,9 @@ def __init__(self, *, data_augmentation, hflip_prob=0.5, mean=(123.0, 117.0, 104
         elif data_augmentation == "lsj":
             self.transforms = T.Compose(
                 [
-                    T.ScaleJitter(target_size=(1024, 1024)),
-                    T.FixedSizeCrop(size=(1024, 1024), fill=mean),
                     T.RandomHorizontalFlip(p=hflip_prob),
+                    T.ScaleJitter(target_size=(1024, 1024)),
+                    T.FixedSizeCrop(size=(1024, 1024), fill=0),
                     T.PILToTensor(),
                     T.ConvertImageDtype(torch.float),
                 ]

references/detection/train.py

@@ -126,6 +126,12 @@ def get_args_parser(add_help=True):
     parser.add_argument(
         "--data-augmentation", default="hflip", type=str, help="data augmentation policy (default: hflip)"
     )
+    parser.add_argument(
+        "--image-min-size", default=800, type=int, help="resize images so that the smallest side is equal to this"
+    )
+    parser.add_argument(
+        "--image-max-size", default=1333, type=int, help="resize images so that the largest side is less than this"
+    )
     parser.add_argument(
         "--sync-bn",
         dest="sync_bn",
@@ -210,7 +216,11 @@ def main(args):
     )
 
     print("Creating model")
-    kwargs = {"trainable_backbone_layers": args.trainable_backbone_layers}
+    kwargs = {
+        "trainable_backbone_layers": args.trainable_backbone_layers,
+        "min_size": args.image_min_size,
+        "max_size": args.image_max_size,
+    }
     if args.data_augmentation in ["multiscale", "lsj"]:
         kwargs["_skip_resize"] = True
     if "rcnn" in args.model:

torchvision/ops/feature_pyramid_network.py

@@ -1,6 +1,7 @@
 from collections import OrderedDict
 from typing import Callable, Dict, List, Optional, Tuple
 
+import torch
 import torch.nn.functional as F
 from torch import nn, Tensor