diff --git a/configs/train.yaml b/configs/train.yaml index 1c52222..7735c66 100644 --- a/configs/train.yaml +++ b/configs/train.yaml @@ -8,9 +8,9 @@ defaults: # these are the sub-config files for specific details, noee for the mo # Training Configuration training_config: - batch_size: 8 - num_workers: 16 - epochs: 200 + batch_size: 4 + num_workers: 8 + epochs: 1000 accum_steps: 1 # Simulate batch size * accum_steps optimizer_momentum: 0.9 optimizer_weight_decay: 5e-4 diff --git a/src/datasets/aeroscapes.py b/src/datasets/aeroscapes.py index 81c9abf..73256a1 100644 --- a/src/datasets/aeroscapes.py +++ b/src/datasets/aeroscapes.py @@ -25,6 +25,7 @@ RandomScale, RandomTranslate, RandomVerticalFlip, + ResizeIfLarger, ) # Same knobs/semantics as src/datasets/uavid.py's DEFAULT_AUGMENTATION — @@ -174,6 +175,12 @@ def __init__( self.trans_train = ( Compose( [ + # No-op at AeroScapes' native 1280x720 with the default + # cropsize, but bounds the working resolution before the + # expensive geometric ops below regardless of cropsize — + # see ResizeIfLarger's docstring (uavid.py/vdd.py hit + # this for real, at their much larger native resolutions). + ResizeIfLarger(max_size=2 * max(self.cropsize)), RandomHorizontalFlip(p=float(self.aug["fliplr"])), RandomVerticalFlip(p=float(self.aug["flipud"])), RandomTranslate( diff --git a/src/datasets/transform.py b/src/datasets/transform.py index d05470f..2e0a550 100644 --- a/src/datasets/transform.py +++ b/src/datasets/transform.py @@ -26,6 +26,42 @@ def __call__(self, im_lb): return im_lb +class ResizeIfLarger(object): + """Downsize (never upsize) so the image's longer side is at most + ``max_size``, preserving aspect ratio. A no-op when the image already fits. + + Purpose: the geometric transforms that typically run before ``RandomCrop`` + in this pipeline (``RandomTranslate``, ``RandomRotate``, ``RandomScale``) + are CPU-bound PIL/numpy operations whose cost scales with the *native* + image resolution, not the eventual training crop size. For datasets with + very large native images (e.g. VDD's 4000x3000, ~11x a 1024x1024 crop), + this makes CPU-side augmentation the training bottleneck — the GPU sits + idle waiting for augmented batches (low, bursty utilization). Placing this + as the *first* step in the pipeline bounds per-sample CPU cost regardless + of native resolution, with no effect on the final crop (still exactly + ``cropsize``) or on crop-position randomness, as long as ``max_size`` is + kept comfortably larger than ``cropsize``. + """ + + def __init__(self, max_size, interp_image=Image.BILINEAR, interp_label=Image.NEAREST): + self.max_size = max_size + self.interp_image = interp_image + self.interp_label = interp_label + + def __call__(self, im_lb): + im = im_lb["im"] + lb = im_lb["lb"] + w, h = im.size + longest = max(w, h) + if longest <= self.max_size: + return im_lb + scale = self.max_size / longest + new_w, new_h = max(1, round(w * scale)), max(1, round(h * scale)) + im_lb["im"] = im.resize((new_w, new_h), self.interp_image) + im_lb["lb"] = lb.resize((new_w, new_h), self.interp_label) + return im_lb + + class RandomScale(object): """Random resize, either from a discrete list of scale factors (default) or a continuous range — pass ``continuous=True`` with diff --git a/src/datasets/uavid.py b/src/datasets/uavid.py index 5b10676..d7e5bf5 100644 --- a/src/datasets/uavid.py +++ b/src/datasets/uavid.py @@ -25,6 +25,7 @@ RandomScale, RandomTranslate, RandomVerticalFlip, + ResizeIfLarger, ) # Mirrors configs/yolo/uavid_train.yaml / configs/train_yolo.yaml's @@ -191,6 +192,12 @@ def __init__( self.trans_train = ( Compose( [ + # UAVid's native images run up to 4096x2160 (~8.85MP) — + # cap the working resolution before the expensive + # geometric ops below, or CPU-side augmentation becomes + # the training bottleneck (GPU starves waiting on the + # DataLoader). See ResizeIfLarger's docstring. + ResizeIfLarger(max_size=2 * max(self.cropsize)), RandomHorizontalFlip(p=float(self.aug["fliplr"])), RandomVerticalFlip(p=float(self.aug["flipud"])), RandomTranslate( diff --git a/src/datasets/vdd.py b/src/datasets/vdd.py index be66558..f3810ab 100644 --- a/src/datasets/vdd.py +++ b/src/datasets/vdd.py @@ -25,6 +25,7 @@ RandomScale, RandomTranslate, RandomVerticalFlip, + ResizeIfLarger, ) # Same knobs/semantics as src/datasets/uavid.py's DEFAULT_AUGMENTATION — VDD @@ -174,6 +175,13 @@ def __init__( self.trans_train = ( Compose( [ + # VDD's native images are 4000x3000 (~11x a 1024x1024 + # crop) — cap the working resolution before the + # expensive geometric ops below, or CPU-side + # augmentation becomes the training bottleneck (GPU + # starves waiting on the DataLoader). See + # ResizeIfLarger's docstring. + ResizeIfLarger(max_size=2 * max(self.cropsize)), RandomHorizontalFlip(p=float(self.aug["fliplr"])), RandomVerticalFlip(p=float(self.aug["flipud"])), RandomTranslate( diff --git a/tests/unit/test_transforms.py b/tests/unit/test_transforms.py index debf5f0..cba8ec9 100644 --- a/tests/unit/test_transforms.py +++ b/tests/unit/test_transforms.py @@ -17,6 +17,7 @@ RandomScale, RandomTranslate, RandomVerticalFlip, + ResizeIfLarger, ) @@ -384,3 +385,49 @@ def test_large_gains_change_pixel_values(self): after = np.array(result["im"]) assert not np.array_equal(before, after) + + +class TestResizeIfLarger: + def test_noop_when_already_within_max_size(self): + img = Image.new("RGB", (256, 256), color=(1, 2, 3)) + lb = Image.new("L", (256, 256), color=1) + transform = ResizeIfLarger(max_size=1024) + result = transform({"im": img, "lb": lb}) + assert result["im"].size == (256, 256) + assert result["im"] is img # untouched, not even a copy + + def test_noop_when_exactly_at_max_size(self): + img = Image.new("RGB", (1024, 512), color=(1, 2, 3)) + lb = Image.new("L", (1024, 512), color=1) + transform = ResizeIfLarger(max_size=1024) + result = transform({"im": img, "lb": lb}) + assert result["im"].size == (1024, 512) + + def test_downsizes_larger_image_preserving_aspect_ratio(self): + img = Image.new("RGB", (4000, 3000), color=(1, 2, 3)) + lb = Image.new("L", (4000, 3000), color=1) + transform = ResizeIfLarger(max_size=2000) + result = transform({"im": img, "lb": lb}) + assert result["im"].size == (2000, 1500) + assert result["lb"].size == (2000, 1500) + + def test_never_upsizes(self): + img = Image.new("RGB", (100, 80), color=(1, 2, 3)) + lb = Image.new("L", (100, 80), color=1) + transform = ResizeIfLarger(max_size=2000) + result = transform({"im": img, "lb": lb}) + assert result["im"].size == (100, 80) + + def test_caps_on_the_longer_dimension_for_portrait_images(self): + img = Image.new("RGB", (1000, 4000), color=(1, 2, 3)) + lb = Image.new("L", (1000, 4000), color=1) + transform = ResizeIfLarger(max_size=2000) + result = transform({"im": img, "lb": lb}) + assert result["im"].size == (500, 2000) + + def test_label_mode_preserved_after_downsize(self): + img = Image.new("RGB", (4000, 3000), color=(1, 2, 3)) + lb = Image.new("L", (4000, 3000), color=1) + transform = ResizeIfLarger(max_size=2000) + result = transform({"im": img, "lb": lb}) + assert result["lb"].mode == "L"