Upload folder using huggingface_hub

.gitattributes

@@ -33,3 +33,10 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+examples/acoustic_guitar.jpg filter=lfs diff=lfs merge=lfs -text
+examples/chameleon.jpg filter=lfs diff=lfs merge=lfs -text
+examples/library_interior.jpg filter=lfs diff=lfs merge=lfs -text
+examples/monstera_plant.jpg filter=lfs diff=lfs merge=lfs -text
+examples/pizza_board.jpg filter=lfs diff=lfs merge=lfs -text
+examples/red_fox.jpg filter=lfs diff=lfs merge=lfs -text
+examples/spiral_staircase.jpg filter=lfs diff=lfs merge=lfs -text

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright 2026 Kiel University
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

LICENSE-timm.txt

@@ -0,0 +1,201 @@
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "{}"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright 2019 Ross Wightman
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

NOTICE

@@ -0,0 +1,20 @@
+=======================================================================
+pytorch-image-models's Apache 2.0 license
+=======================================================================
+We modified and utilize the pytorch-image-models implementation from
+https://github.com/huggingface/pytorch-image-models. The license application is as
+follows:
+
+Copyright 2019 Ross Wightman
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.

README.md

@@ -1,13 +1,70 @@
 ---
-title: Progresvit Classify
-emoji: 💻
-colorFrom: blue
-colorTo: green
+title: ProgResViT
+emoji: 🪜
+colorFrom: indigo
+colorTo: purple
 sdk: gradio
 sdk_version: 6.26.0
-python_version: '3.12'
 app_file: app.py
-pinned: false
+short_description: Adaptive-compute ViT that classifies in progressive rounds
+python_version: "3.12"
+startup_duration_timeout: 30m
+license: mit
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# ProgResViT: Progressive Resolution and Width for Adaptive Vision Transformers
+
+Interactive ImageNet-1K demo of [ProgResViT](https://huggingface.co/papers/2609.03216)
+(arXiv:2609.03216, Kiel University).
+
+ProgResViT performs inference **progressively**. Round 1 processes a low-resolution
+image with a narrow subnetwork (3 of 6 attention heads). If the round-1 prediction is
+confident enough — measured by the entropy of its top-10 softmax — inference stops
+there. Otherwise the model recycles the round-1 tokens and refines the prediction at a
+higher input resolution with the full-width subnetwork. All rounds share a single
+backbone, conditioned by **Progress-Conditioned Soft Gating (PSG)**.
+
+The demo exposes that mechanism directly: it runs both rounds, shows each round's top-5
+prediction, and reports which round the routing threshold would have stopped at, along
+with the GMACs saved.
+
+## Checkpoints
+
+All four released DeiT-S checkpoints are available in the dropdown:
+
+| Resolution schedule | Training | Top-1 | GMACs (full path) |
+|---|---|---:|---:|
+| 160 → 384 | KD | 84.90% | 16.152 |
+| 160 → 384 | standard | 83.70% | 16.152 |
+| 192 → 240 | KD | 83.80% | 6.267 |
+| 192 → 240 | standard | 82.21% | 6.267 |
+
+Weights: [NCPS on the Hub](https://huggingface.co/NCPS). Default routing thresholds are
+the authors' reported operating points (≤0.03 pp top-1 drop).
+
+## Implementation notes
+
+- The ProgResViT model code is the authors' vendored `timm` fork, copied verbatim from
+  [ds-kiel/ProgResViT](https://github.com/ds-kiel/ProgResViT) (MIT; `NOTICE` and
+  `LICENSE-timm.txt` retained).
+- Preprocessing matches `validate.py` upstream: bicubic resize with `crop_pct=0.9`,
+  center crop to the checkpoint's eval resolution, ImageNet mean/std.
+- Rounds are run with `model._forward_stage(...)` exactly as the upstream evaluator
+  does, so both rounds are always computed and the routing decision is reported rather
+  than short-circuited — that is what makes the trade-off visible.
+- GMACs figures are the authors' measured values from `results/RESULTS.md`.
+
+## Credits
+
+Example photographs come from
+[linoyts/repo-to-space-example-inputs](https://huggingface.co/datasets/linoyts/repo-to-space-example-inputs).
+
+## Citation
+
+```bibtex
+@article{progresvit2026,
+  title  = {ProgResViT: Progressive Resolution and Width for Adaptive Vision Transformers},
+  year   = {2026},
+  eprint = {2609.03216}
+}
+```

app.py

@@ -0,0 +1,280 @@
+"""ProgResViT — progressive-resolution / progressive-width adaptive ViT.
+
+Interactive ImageNet-1K classification demo that exposes the paper's
+input-adaptive routing: round 1 runs a narrow subnetwork on a low-resolution
+image, and only uncertain images continue to round 2 at higher resolution and
+wider width.
+
+Paper: https://huggingface.co/papers/2609.03216
+Code:  https://github.com/ds-kiel/ProgResViT
+"""
+
+import json
+import os
+import time
+
+import spaces  # must precede torch
+import torch
+import gradio as gr
+from PIL import Image
+from huggingface_hub import hf_hub_download
+from safetensors.torch import load_file
+
+from timm.data.transforms_factory import create_transform
+from timm.models import create_model
+
+# ---------------------------------------------------------------------------
+# Model registry
+# ---------------------------------------------------------------------------
+# GMACs come from the authors' measured sweeps (results/RESULTS.md in the
+# upstream repo): the threshold=0 row is the full two-round cost, the
+# threshold=10 row (every image exits after round 1) is the round-1 cost.
+VARIANTS = {
+    "160 → 384 · KD (84.9% top-1)": {
+        "repo": "NCPS/progresvit-deit-s-160-384-kd-imagenet1k",
+        "sizes": (160, 384),
+        "gmacs": (0.615, 16.152),
+        "top1": (73.940, 84.894),
+        "amp": True,
+        "threshold": 0.226,
+    },
+    "160 → 384 (83.7% top-1)": {
+        "repo": "NCPS/progresvit-deit-s-160-384-imagenet1k",
+        "sizes": (160, 384),
+        "gmacs": (0.615, 16.152),
+        "top1": (70.616, 83.714),
+        "amp": False,
+        "threshold": 0.267,
+    },
+    "192 → 240 · KD (83.8% top-1)": {
+        "repo": "NCPS/progresvit-deit-s-192-240-kd-imagenet1k",
+        "sizes": (192, 240),
+        "gmacs": (0.912, 6.267),
+        "top1": (76.018, 83.794),
+        "amp": False,
+        "threshold": 0.209,
+    },
+    "192 → 240 (82.2% top-1)": {
+        "repo": "NCPS/progresvit-deit-s-192-240-imagenet1k",
+        "sizes": (192, 240),
+        "gmacs": (0.912, 6.267),
+        "top1": (73.238, 82.202),
+        "amp": False,
+        "threshold": 0.356,
+    },
+}
+
+DEFAULT_VARIANT = "160 → 384 · KD (84.9% top-1)"
+DEFAULT_THRESHOLD = VARIANTS[DEFAULT_VARIANT]["threshold"]
+PROGRESS_STAGES = (3, 6)  # attention heads active in round 1 / round 2
+CACHE_VERSION = 1
+
+with open(os.path.join(os.path.dirname(__file__), "imagenet_classes.json")) as f:
+    IMAGENET_CLASSES = [json.load(f)[str(i)] for i in range(1000)]
+
+MODELS = {}
+TRANSFORMS = {}
+CROPS = {}
+
+for _name, _spec in VARIANTS.items():
+    _cfg = json.load(open(hf_hub_download(_spec["repo"], "config.json")))
+    _model = create_model(
+        "progresvit",
+        pretrained=False,
+        num_classes=_cfg["num_classes"],
+        **_cfg["model_args"],
+    )
+    _state = load_file(hf_hub_download(_spec["repo"], "model.safetensors"))
+    _model.load_state_dict(_state, strict=True)
+    _pc = _cfg["pretrained_cfg"]
+    TRANSFORMS[_name] = create_transform(
+        input_size=tuple(_pc["input_size"]),
+        is_training=False,
+        interpolation=_pc["interpolation"],
+        mean=tuple(_pc["mean"]),
+        std=tuple(_pc["std"]),
+        crop_pct=_pc["crop_pct"],
+        crop_mode=_pc["crop_mode"],
+        crop_border_pixels=0,
+        use_prefetcher=False,
+    )
+    CROPS[_name] = int(_pc["input_size"][-1])
+    MODELS[_name] = _model.eval().to("cuda")
+    print(f"loaded {_name} from {_spec['repo']} (eval crop {CROPS[_name]})", flush=True)
+
+
+def _topk_dict(logits: torch.Tensor, k: int = 5) -> dict:
+    probs = logits.float().softmax(dim=-1)[0]
+    values, indices = probs.topk(k)
+    return {IMAGENET_CLASSES[int(i)]: float(v) for v, i in zip(values, indices)}
+
+
+@spaces.GPU(duration=30)
+def classify(
+    image: Image.Image,
+    variant: str = DEFAULT_VARIANT,
+    threshold: float = DEFAULT_THRESHOLD,
+) -> tuple:
+    """Classify an image with ProgResViT's progressive, input-adaptive rounds.
+
+    Args:
+        image: input photograph to classify against the 1000 ImageNet-1K classes.
+        variant: which ProgResViT DeiT-S checkpoint to use (resolution schedule
+            and whether it was trained with knowledge distillation).
+        threshold: routing threshold on the round-1 top-10 prediction entropy.
+            The image exits after the cheap first round when its entropy falls
+            below this value; higher values exit more images and save more
+            compute.
+
+    Returns:
+        A tuple of (final top-5 prediction, routing report in markdown,
+        round-1 top-5 prediction, round-2 top-5 prediction).
+    """
+    if image is None:
+        raise gr.Error("Please provide an image.")
+
+    spec = VARIANTS[variant]
+    model = MODELS[variant]
+    sizes = spec["sizes"]
+    g1, g2 = spec["gmacs"]
+
+    x = TRANSFORMS[variant](image.convert("RGB")).unsqueeze(0).to("cuda")
+
+    started = time.perf_counter()
+    with torch.inference_mode():
+        if spec["amp"]:
+            ctx = torch.autocast("cuda", dtype=torch.bfloat16)
+        else:
+            ctx = torch.autocast("cuda", enabled=False)
+        with ctx:
+            tokens1, logits1 = model._forward_stage(
+                x, 0, None, PROGRESS_STAGES, sizes
+            )
+            _, logits2 = model._forward_stage(
+                x, 1, tokens1, PROGRESS_STAGES, sizes
+            )
+        entropy = float(model.entropy(logits1.float())[0, 0])
+    elapsed = time.perf_counter() - started
+
+    exited_early = entropy < threshold
+    final_logits = logits1 if exited_early else logits2
+    used_gmacs = g1 if exited_early else g2
+    saving = 100.0 * (1.0 - used_gmacs / g2)
+
+    round1 = _topk_dict(logits1)
+    round2 = _topk_dict(logits2)
+    final = _topk_dict(final_logits)
+
+    if exited_early:
+        decision = (
+            f"**Exited after round 1.** Entropy `{entropy:.3f}` is below the "
+            f"threshold `{threshold:.3f}`, so the {sizes[1]} px round was skipped."
+        )
+    else:
+        decision = (
+            f"**Continued to round 2.** Entropy `{entropy:.3f}` is at or above the "
+            f"threshold `{threshold:.3f}`, so round 1's tokens were recycled and "
+            f"refined at {sizes[1]} px."
+        )
+
+    report = f"""### Routing
+
+{decision}
+
+| | Round 1 | Round 2 | This image |
+|---|---|---|---|
+| Input resolution | {sizes[0]} px | {sizes[1]} px | **{sizes[0] if exited_early else sizes[1]} px** |
+| Active attention heads | {PROGRESS_STAGES[0]} / 6 | {PROGRESS_STAGES[1]} / 6 | **{PROGRESS_STAGES[0] if exited_early else PROGRESS_STAGES[1]} / 6** |
+| Cumulative GMACs | {g1:.3f} | {g2:.3f} | **{used_gmacs:.3f}** |
+| ImageNet top-1 if always stopped here | {spec['top1'][0]:.2f}% | {spec['top1'][1]:.2f}% | — |
+
+Compute saved versus always running both rounds: **{saving:.1f}%** · inference {elapsed * 1000:.0f} ms
+"""
+    return final, report, round1, round2
+
+
+CSS = """
+#col-container { max-width: 1180px; margin: 0 auto; }
+.dark .gradio-container { color: var(--body-text-color); }
+"""
+
+# Ordered so the first rows tell the story: `red_fox` stays uncertain after round 1
+# (which calls it a kit fox) and gets corrected in round 2, while `acoustic_guitar`
+# is confident enough to exit after the cheap first round.
+EXAMPLES = [
+    ["examples/red_fox.jpg"],
+    ["examples/acoustic_guitar.jpg"],
+    ["examples/husky_dog.jpg"],
+    ["examples/pizza_board.jpg"],
+    ["examples/bird_kingfisher.jpg"],
+    ["examples/chameleon.jpg"],
+    ["examples/hot_air_balloon.jpg"],
+    ["examples/vintage_camera.jpg"],
+    ["examples/library_interior.jpg"],
+    ["examples/spiral_staircase.jpg"],
+    ["examples/monstera_plant.jpg"],
+]
+
+with gr.Blocks(theme=gr.themes.Citrus(), css=CSS) as demo:
+    with gr.Column(elem_id="col-container"):
+        gr.Markdown(
+            """# ProgResViT — adaptive-compute image classification
+
+An input-adaptive Vision Transformer that classifies progressively: round 1 runs a
+**narrow** subnetwork on a **low-resolution** image, and only images whose prediction is
+still uncertain continue to round 2 at **higher resolution** with a **wider** subnetwork,
+reusing the tokens produced in round 1.
+
+[Paper](https://huggingface.co/papers/2609.03216) · [Code](https://github.com/ds-kiel/ProgResViT) · [Checkpoints](https://huggingface.co/NCPS)
+"""
+        )
+        with gr.Row():
+            with gr.Column():
+                image = gr.Image(label="Image", type="pil", height=340)
+                run = gr.Button("Classify", variant="primary")
+                variant = gr.Dropdown(
+                    label="Checkpoint",
+                    choices=list(VARIANTS),
+                    value=DEFAULT_VARIANT,
+                )
+                threshold = gr.Slider(
+                    label="Routing threshold (round-1 entropy)",
+                    minimum=0.0,
+                    maximum=2.0,
+                    step=0.001,
+                    value=DEFAULT_THRESHOLD,
+                    info="0 = always run both rounds · higher = exit more images early",
+                )
+            with gr.Column():
+                final_out = gr.Label(label="Prediction", num_top_classes=5)
+                report_out = gr.Markdown()
+
+        with gr.Accordion("Round-by-round predictions", open=False):
+            with gr.Row():
+                round1_out = gr.Label(label="Round 1 (low-res, narrow)", num_top_classes=5)
+                round2_out = gr.Label(label="Round 2 (high-res, wide)", num_top_classes=5)
+
+        gr.Examples(
+            examples=EXAMPLES,
+            inputs=[image],
+            outputs=[final_out, report_out, round1_out, round2_out],
+            fn=classify,
+            cache_examples=True,
+            cache_mode="lazy",
+            examples_per_page=12,
+        )
+
+    def _sync_threshold(name: str) -> float:
+        """Reset the routing threshold to the checkpoint's reported operating point."""
+        return VARIANTS[name]["threshold"]
+
+    variant.change(_sync_threshold, inputs=variant, outputs=threshold)
+    run.click(
+        classify,
+        inputs=[image, variant, threshold],
+        outputs=[final_out, report_out, round1_out, round2_out],
+        api_name="classify",
+    )
+
+if __name__ == "__main__":
+    demo.launch(mcp_server=True)

imagenet_classes.json

@@ -0,0 +1,1002 @@
+{
+  "0": "tench, Tinca tinca",
+  "1": "goldfish, Carassius auratus",
+  "2": "great white shark, white shark, man-eater, man-eating shark, Carcharodon carcharias",
+  "3": "tiger shark, Galeocerdo cuvieri",
+  "4": "hammerhead, hammerhead shark",
+  "5": "electric ray, crampfish, numbfish, torpedo",
+  "6": "stingray",
+  "7": "cock",
+  "8": "hen",
+  "9": "ostrich, Struthio camelus",
+  "10": "brambling, Fringilla montifringilla",
+  "11": "goldfinch, Carduelis carduelis",
+  "12": "house finch, linnet, Carpodacus mexicanus",
+  "13": "junco, snowbird",
+  "14": "indigo bunting, indigo finch, indigo bird, Passerina cyanea",
+  "15": "robin, American robin, Turdus migratorius",
+  "16": "bulbul",
+  "17": "jay",
+  "18": "magpie",
+  "19": "chickadee",
+  "20": "water ouzel, dipper",
+  "21": "kite",
+  "22": "bald eagle, American eagle, Haliaeetus leucocephalus",
+  "23": "vulture",
+  "24": "great grey owl, great gray owl, Strix nebulosa",
+  "25": "European fire salamander, Salamandra salamandra",
+  "26": "common newt, Triturus vulgaris",
+  "27": "eft",
+  "28": "spotted salamander, Ambystoma maculatum",
+  "29": "axolotl, mud puppy, Ambystoma mexicanum",
+  "30": "bullfrog, Rana catesbeiana",
+  "31": "tree frog, tree-frog",
+  "32": "tailed frog, bell toad, ribbed toad, tailed toad, Ascaphus trui",
+  "33": "loggerhead, loggerhead turtle, Caretta caretta",
+  "34": "leatherback turtle, leatherback, leathery turtle, Dermochelys coriacea",
+  "35": "mud turtle",
+  "36": "terrapin",
+  "37": "box turtle, box tortoise",
+  "38": "banded gecko",
+  "39": "common iguana, iguana, Iguana iguana",
+  "40": "American chameleon, anole, Anolis carolinensis",
+  "41": "whiptail, whiptail lizard",
+  "42": "agama",
+  "43": "frilled lizard, Chlamydosaurus kingi",
+  "44": "alligator lizard",
+  "45": "Gila monster, Heloderma suspectum",
+  "46": "green lizard, Lacerta viridis",
+  "47": "African chameleon, Chamaeleo chamaeleon",
+  "48": "Komodo dragon, Komodo lizard, dragon lizard, giant lizard, Varanus komodoensis",
+  "49": "African crocodile, Nile crocodile, Crocodylus niloticus",
+  "50": "American alligator, Alligator mississipiensis",
+  "51": "triceratops",
+  "52": "thunder snake, worm snake, Carphophis amoenus",
+  "53": "ringneck snake, ring-necked snake, ring snake",
+  "54": "hognose snake, puff adder, sand viper",
+  "55": "green snake, grass snake",
+  "56": "king snake, kingsnake",
+  "57": "garter snake, grass snake",
+  "58": "water snake",
+  "59": "vine snake",
+  "60": "night snake, Hypsiglena torquata",
+  "61": "boa constrictor, Constrictor constrictor",
+  "62": "rock python, rock snake, Python sebae",
+  "63": "Indian cobra, Naja naja",
+  "64": "green mamba",
+  "65": "sea snake",
+  "66": "horned viper, cerastes, sand viper, horned asp, Cerastes cornutus",
+  "67": "diamondback, diamondback rattlesnake, Crotalus adamanteus",
+  "68": "sidewinder, horned rattlesnake, Crotalus cerastes",
+  "69": "trilobite",
+  "70": "harvestman, daddy longlegs, Phalangium opilio",
+  "71": "scorpion",
+  "72": "black and gold garden spider, Argiope aurantia",
+  "73": "barn spider, Araneus cavaticus",
+  "74": "garden spider, Aranea diademata",
+  "75": "black widow, Latrodectus mactans",
+  "76": "tarantula",
+  "77": "wolf spider, hunting spider",
+  "78": "tick",
+  "79": "centipede",
+  "80": "black grouse",
+  "81": "ptarmigan",
+  "82": "ruffed grouse, partridge, Bonasa umbellus",
+  "83": "prairie chicken, prairie grouse, prairie fowl",
+  "84": "peacock",
+  "85": "quail",
+  "86": "partridge",
+  "87": "African grey, African gray, Psittacus erithacus",
+  "88": "macaw",
+  "89": "sulphur-crested cockatoo, Kakatoe galerita, Cacatua galerita",
+  "90": "lorikeet",
+  "91": "coucal",
+  "92": "bee eater",
+  "93": "hornbill",
+  "94": "hummingbird",
+  "95": "jacamar",
+  "96": "toucan",
+  "97": "drake",
+  "98": "red-breasted merganser, Mergus serrator",
+  "99": "goose",
+  "100": "black swan, Cygnus atratus",
+  "101": "tusker",
+  "102": "echidna, spiny anteater, anteater",
+  "103": "platypus, duckbill, duckbilled platypus, duck-billed platypus, Ornithorhynchus anatinus",
+  "104": "wallaby, brush kangaroo",
+  "105": "koala, koala bear, kangaroo bear, native bear, Phascolarctos cinereus",
+  "106": "wombat",
+  "107": "jellyfish",
+  "108": "sea anemone, anemone",
+  "109": "brain coral",
+  "110": "flatworm, platyhelminth",
+  "111": "nematode, nematode worm, roundworm",
+  "112": "conch",
+  "113": "snail",
+  "114": "slug",
+  "115": "sea slug, nudibranch",
+  "116": "chiton, coat-of-mail shell, sea cradle, polyplacophore",
+  "117": "chambered nautilus, pearly nautilus, nautilus",
+  "118": "Dungeness crab, Cancer magister",
+  "119": "rock crab, Cancer irroratus",
+  "120": "fiddler crab",
+  "121": "king crab, Alaska crab, Alaskan king crab, Alaska king crab, Paralithodes camtschatica",
+  "122": "American lobster, Northern lobster, Maine lobster, Homarus americanus",
+  "123": "spiny lobster, langouste, rock lobster, crawfish, crayfish, sea crawfish",
+  "124": "crayfish, crawfish, crawdad, crawdaddy",
+  "125": "hermit crab",
+  "126": "isopod",
+  "127": "white stork, Ciconia ciconia",
+  "128": "black stork, Ciconia nigra",
+  "129": "spoonbill",
+  "130": "flamingo",
+  "131": "little blue heron, Egretta caerulea",
+  "132": "American egret, great white heron, Egretta albus",
+  "133": "bittern",
+  "134": "crane",
+  "135": "limpkin, Aramus pictus",
+  "136": "European gallinule, Porphyrio porphyrio",
+  "137": "American coot, marsh hen, mud hen, water hen, Fulica americana",
+  "138": "bustard",
+  "139": "ruddy turnstone, Arenaria interpres",
+  "140": "red-backed sandpiper, dunlin, Erolia alpina",
+  "141": "redshank, Tringa totanus",
+  "142": "dowitcher",
+  "143": "oystercatcher, oyster catcher",
+  "144": "pelican",
+  "145": "king penguin, Aptenodytes patagonica",
+  "146": "albatross, mollymawk",
+  "147": "grey whale, gray whale, devilfish, Eschrichtius gibbosus, Eschrichtius robustus",
+  "148": "killer whale, killer, orca, grampus, sea wolf, Orcinus orca",
+  "149": "dugong, Dugong dugon",
+  "150": "sea lion",
+  "151": "Chihuahua",
+  "152": "Japanese spaniel",
+  "153": "Maltese dog, Maltese terrier, Maltese",
+  "154": "Pekinese, Pekingese, Peke",
+  "155": "Shih-Tzu",
+  "156": "Blenheim spaniel",
+  "157": "papillon",
+  "158": "toy terrier",
+  "159": "Rhodesian ridgeback",
+  "160": "Afghan hound, Afghan",
+  "161": "basset, basset hound",
+  "162": "beagle",
+  "163": "bloodhound, sleuthhound",
+  "164": "bluetick",
+  "165": "black-and-tan coonhound",
+  "166": "Walker hound, Walker foxhound",
+  "167": "English foxhound",
+  "168": "redbone",
+  "169": "borzoi, Russian wolfhound",
+  "170": "Irish wolfhound",
+  "171": "Italian greyhound",
+  "172": "whippet",
+  "173": "Ibizan hound, Ibizan Podenco",
+  "174": "Norwegian elkhound, elkhound",
+  "175": "otterhound, otter hound",
+  "176": "Saluki, gazelle hound",
+  "177": "Scottish deerhound, deerhound",
+  "178": "Weimaraner",
+  "179": "Staffordshire bullterrier, Staffordshire bull terrier",
+  "180": "American Staffordshire terrier, Staffordshire terrier, American pit bull terrier, pit bull terrier",
+  "181": "Bedlington terrier",
+  "182": "Border terrier",
+  "183": "Kerry blue terrier",
+  "184": "Irish terrier",
+  "185": "Norfolk terrier",
+  "186": "Norwich terrier",
+  "187": "Yorkshire terrier",
+  "188": "wire-haired fox terrier",
+  "189": "Lakeland terrier",
+  "190": "Sealyham terrier, Sealyham",
+  "191": "Airedale, Airedale terrier",
+  "192": "cairn, cairn terrier",
+  "193": "Australian terrier",
+  "194": "Dandie Dinmont, Dandie Dinmont terrier",
+  "195": "Boston bull, Boston terrier",
+  "196": "miniature schnauzer",
+  "197": "giant schnauzer",
+  "198": "standard schnauzer",
+  "199": "Scotch terrier, Scottish terrier, Scottie",
+  "200": "Tibetan terrier, chrysanthemum dog",
+  "201": "silky terrier, Sydney silky",
+  "202": "soft-coated wheaten terrier",
+  "203": "West Highland white terrier",
+  "204": "Lhasa, Lhasa apso",
+  "205": "flat-coated retriever",
+  "206": "curly-coated retriever",
+  "207": "golden retriever",
+  "208": "Labrador retriever",
+  "209": "Chesapeake Bay retriever",
+  "210": "German short-haired pointer",
+  "211": "vizsla, Hungarian pointer",
+  "212": "English setter",
+  "213": "Irish setter, red setter",
+  "214": "Gordon setter",
+  "215": "Brittany spaniel",
+  "216": "clumber, clumber spaniel",
+  "217": "English springer, English springer spaniel",
+  "218": "Welsh springer spaniel",
+  "219": "cocker spaniel, English cocker spaniel, cocker",
+  "220": "Sussex spaniel",
+  "221": "Irish water spaniel",
+  "222": "kuvasz",
+  "223": "schipperke",
+  "224": "groenendael",
+  "225": "malinois",
+  "226": "briard",
+  "227": "kelpie",
+  "228": "komondor",
+  "229": "Old English sheepdog, bobtail",
+  "230": "Shetland sheepdog, Shetland sheep dog, Shetland",
+  "231": "collie",
+  "232": "Border collie",
+  "233": "Bouvier des Flandres, Bouviers des Flandres",
+  "234": "Rottweiler",
+  "235": "German shepherd, German shepherd dog, German police dog, alsatian",
+  "236": "Doberman, Doberman pinscher",
+  "237": "miniature pinscher",
+  "238": "Greater Swiss Mountain dog",
+  "239": "Bernese mountain dog",
+  "240": "Appenzeller",
+  "241": "EntleBucher",
+  "242": "boxer",
+  "243": "bull mastiff",
+  "244": "Tibetan mastiff",
+  "245": "French bulldog",
+  "246": "Great Dane",
+  "247": "Saint Bernard, St Bernard",
+  "248": "Eskimo dog, husky",
+  "249": "malamute, malemute, Alaskan malamute",
+  "250": "Siberian husky",
+  "251": "dalmatian, coach dog, carriage dog",
+  "252": "affenpinscher, monkey pinscher, monkey dog",
+  "253": "basenji",
+  "254": "pug, pug-dog",
+  "255": "Leonberg",
+  "256": "Newfoundland, Newfoundland dog",
+  "257": "Great Pyrenees",
+  "258": "Samoyed, Samoyede",
+  "259": "Pomeranian",
+  "260": "chow, chow chow",
+  "261": "keeshond",
+  "262": "Brabancon griffon",
+  "263": "Pembroke, Pembroke Welsh corgi",
+  "264": "Cardigan, Cardigan Welsh corgi",
+  "265": "toy poodle",
+  "266": "miniature poodle",
+  "267": "standard poodle",
+  "268": "Mexican hairless",
+  "269": "timber wolf, grey wolf, gray wolf, Canis lupus",
+  "270": "white wolf, Arctic wolf, Canis lupus tundrarum",
+  "271": "red wolf, maned wolf, Canis rufus, Canis niger",
+  "272": "coyote, prairie wolf, brush wolf, Canis latrans",
+  "273": "dingo, warrigal, warragal, Canis dingo",
+  "274": "dhole, Cuon alpinus",
+  "275": "African hunting dog, hyena dog, Cape hunting dog, Lycaon pictus",
+  "276": "hyena, hyaena",
+  "277": "red fox, Vulpes vulpes",
+  "278": "kit fox, Vulpes macrotis",
+  "279": "Arctic fox, white fox, Alopex lagopus",
+  "280": "grey fox, gray fox, Urocyon cinereoargenteus",
+  "281": "tabby, tabby cat",
+  "282": "tiger cat",
+  "283": "Persian cat",
+  "284": "Siamese cat, Siamese",
+  "285": "Egyptian cat",
+  "286": "cougar, puma, catamount, mountain lion, painter, panther, Felis concolor",
+  "287": "lynx, catamount",
+  "288": "leopard, Panthera pardus",
+  "289": "snow leopard, ounce, Panthera uncia",
+  "290": "jaguar, panther, Panthera onca, Felis onca",
+  "291": "lion, king of beasts, Panthera leo",
+  "292": "tiger, Panthera tigris",
+  "293": "cheetah, chetah, Acinonyx jubatus",
+  "294": "brown bear, bruin, Ursus arctos",
+  "295": "American black bear, black bear, Ursus americanus, Euarctos americanus",
+  "296": "ice bear, polar bear, Ursus Maritimus, Thalarctos maritimus",
+  "297": "sloth bear, Melursus ursinus, Ursus ursinus",
+  "298": "mongoose",
+  "299": "meerkat, mierkat",
+  "300": "tiger beetle",
+  "301": "ladybug, ladybeetle, lady beetle, ladybird, ladybird beetle",
+  "302": "ground beetle, carabid beetle",
+  "303": "long-horned beetle, longicorn, longicorn beetle",
+  "304": "leaf beetle, chrysomelid",
+  "305": "dung beetle",
+  "306": "rhinoceros beetle",
+  "307": "weevil",
+  "308": "fly",
+  "309": "bee",
+  "310": "ant, emmet, pismire",
+  "311": "grasshopper, hopper",
+  "312": "cricket",
+  "313": "walking stick, walkingstick, stick insect",
+  "314": "cockroach, roach",
+  "315": "mantis, mantid",
+  "316": "cicada, cicala",
+  "317": "leafhopper",
+  "318": "lacewing, lacewing fly",
+  "319": "dragonfly, darning needle, devil's darning needle, sewing needle, snake feeder, snake doctor, mosquito hawk, skeeter hawk",
+  "320": "damselfly",
+  "321": "admiral",
+  "322": "ringlet, ringlet butterfly",
+  "323": "monarch, monarch butterfly, milkweed butterfly, Danaus plexippus",
+  "324": "cabbage butterfly",
+  "325": "sulphur butterfly, sulfur butterfly",
+  "326": "lycaenid, lycaenid butterfly",
+  "327": "starfish, sea star",
+  "328": "sea urchin",
+  "329": "sea cucumber, holothurian",
+  "330": "wood rabbit, cottontail, cottontail rabbit",
+  "331": "hare",
+  "332": "Angora, Angora rabbit",
+  "333": "hamster",
+  "334": "porcupine, hedgehog",
+  "335": "fox squirrel, eastern fox squirrel, Sciurus niger",
+  "336": "marmot",
+  "337": "beaver",
+  "338": "guinea pig, Cavia cobaya",
+  "339": "sorrel",
+  "340": "zebra",
+  "341": "hog, pig, grunter, squealer, Sus scrofa",
+  "342": "wild boar, boar, Sus scrofa",
+  "343": "warthog",
+  "344": "hippopotamus, hippo, river horse, Hippopotamus amphibius",
+  "345": "ox",
+  "346": "water buffalo, water ox, Asiatic buffalo, Bubalus bubalis",
+  "347": "bison",
+  "348": "ram, tup",
+  "349": "bighorn, bighorn sheep, cimarron, Rocky Mountain bighorn, Rocky Mountain sheep, Ovis canadensis",
+  "350": "ibex, Capra ibex",
+  "351": "hartebeest",
+  "352": "impala, Aepyceros melampus",
+  "353": "gazelle",
+  "354": "Arabian camel, dromedary, Camelus dromedarius",
+  "355": "llama",
+  "356": "weasel",
+  "357": "mink",
+  "358": "polecat, fitch, foulmart, foumart, Mustela putorius",
+  "359": "black-footed ferret, ferret, Mustela nigripes",
+  "360": "otter",
+  "361": "skunk, polecat, wood pussy",
+  "362": "badger",
+  "363": "armadillo",
+  "364": "three-toed sloth, ai, Bradypus tridactylus",
+  "365": "orangutan, orang, orangutang, Pongo pygmaeus",
+  "366": "gorilla, Gorilla gorilla",
+  "367": "chimpanzee, chimp, Pan troglodytes",
+  "368": "gibbon, Hylobates lar",
+  "369": "siamang, Hylobates syndactylus, Symphalangus syndactylus",
+  "370": "guenon, guenon monkey",
+  "371": "patas, hussar monkey, Erythrocebus patas",
+  "372": "baboon",
+  "373": "macaque",
+  "374": "langur",
+  "375": "colobus, colobus monkey",
+  "376": "proboscis monkey, Nasalis larvatus",
+  "377": "marmoset",
+  "378": "capuchin, ringtail, Cebus capucinus",
+  "379": "howler monkey, howler",
+  "380": "titi, titi monkey",
+  "381": "spider monkey, Ateles geoffroyi",
+  "382": "squirrel monkey, Saimiri sciureus",
+  "383": "Madagascar cat, ring-tailed lemur, Lemur catta",
+  "384": "indri, indris, Indri indri, Indri brevicaudatus",
+  "385": "Indian elephant, Elephas maximus",
+  "386": "African elephant, Loxodonta africana",
+  "387": "lesser panda, red panda, panda, bear cat, cat bear, Ailurus fulgens",
+  "388": "giant panda, panda, panda bear, coon bear, Ailuropoda melanoleuca",
+  "389": "barracouta, snoek",
+  "390": "eel",
+  "391": "coho, cohoe, coho salmon, blue jack, silver salmon, Oncorhynchus kisutch",
+  "392": "rock beauty, Holocanthus tricolor",
+  "393": "anemone fish",
+  "394": "sturgeon",
+  "395": "gar, garfish, garpike, billfish, Lepisosteus osseus",
+  "396": "lionfish",
+  "397": "puffer, pufferfish, blowfish, globefish",
+  "398": "abacus",
+  "399": "abaya",
+  "400": "academic gown, academic robe, judge's robe",
+  "401": "accordion, piano accordion, squeeze box",
+  "402": "acoustic guitar",
+  "403": "aircraft carrier, carrier, flattop, attack aircraft carrier",
+  "404": "airliner",
+  "405": "airship, dirigible",
+  "406": "altar",
+  "407": "ambulance",
+  "408": "amphibian, amphibious vehicle",
+  "409": "analog clock",
+  "410": "apiary, bee house",
+  "411": "apron",
+  "412": "ashcan, trash can, garbage can, wastebin, ash bin, ash-bin, ashbin, dustbin, trash barrel, trash bin",
+  "413": "assault rifle, assault gun",
+  "414": "backpack, back pack, knapsack, packsack, rucksack, haversack",
+  "415": "bakery, bakeshop, bakehouse",
+  "416": "balance beam, beam",
+  "417": "balloon",
+  "418": "ballpoint, ballpoint pen, ballpen, Biro",
+  "419": "Band Aid",
+  "420": "banjo",
+  "421": "bannister, banister, balustrade, balusters, handrail",
+  "422": "barbell",
+  "423": "barber chair",
+  "424": "barbershop",
+  "425": "barn",
+  "426": "barometer",
+  "427": "barrel, cask",
+  "428": "barrow, garden cart, lawn cart, wheelbarrow",
+  "429": "baseball",
+  "430": "basketball",
+  "431": "bassinet",
+  "432": "bassoon",
+  "433": "bathing cap, swimming cap",
+  "434": "bath towel",
+  "435": "bathtub, bathing tub, bath, tub",
+  "436": "beach wagon, station wagon, wagon, estate car, beach waggon, station waggon, waggon",
+  "437": "beacon, lighthouse, beacon light, pharos",
+  "438": "beaker",
+  "439": "bearskin, busby, shako",
+  "440": "beer bottle",
+  "441": "beer glass",
+  "442": "bell cote, bell cot",
+  "443": "bib",
+  "444": "bicycle-built-for-two, tandem bicycle, tandem",
+  "445": "bikini, two-piece",
+  "446": "binder, ring-binder",
+  "447": "binoculars, field glasses, opera glasses",
+  "448": "birdhouse",
+  "449": "boathouse",
+  "450": "bobsled, bobsleigh, bob",
+  "451": "bolo tie, bolo, bola tie, bola",
+  "452": "bonnet, poke bonnet",
+  "453": "bookcase",
+  "454": "bookshop, bookstore, bookstall",
+  "455": "bottlecap",
+  "456": "bow",
+  "457": "bow tie, bow-tie, bowtie",
+  "458": "brass, memorial tablet, plaque",
+  "459": "brassiere, bra, bandeau",
+  "460": "breakwater, groin, groyne, mole, bulwark, seawall, jetty",
+  "461": "breastplate, aegis, egis",
+  "462": "broom",
+  "463": "bucket, pail",
+  "464": "buckle",
+  "465": "bulletproof vest",
+  "466": "bullet train, bullet",
+  "467": "butcher shop, meat market",
+  "468": "cab, hack, taxi, taxicab",
+  "469": "caldron, cauldron",
+  "470": "candle, taper, wax light",
+  "471": "cannon",
+  "472": "canoe",
+  "473": "can opener, tin opener",
+  "474": "cardigan",
+  "475": "car mirror",
+  "476": "carousel, carrousel, merry-go-round, roundabout, whirligig",
+  "477": "carpenter's kit, tool kit",
+  "478": "carton",
+  "479": "car wheel",
+  "480": "cash machine, cash dispenser, automated teller machine, automatic teller machine, automated teller, automatic teller, ATM",
+  "481": "cassette",
+  "482": "cassette player",
+  "483": "castle",
+  "484": "catamaran",
+  "485": "CD player",
+  "486": "cello, violoncello",
+  "487": "cellular telephone, cellular phone, cellphone, cell, mobile phone",
+  "488": "chain",
+  "489": "chainlink fence",
+  "490": "chain mail, ring mail, mail, chain armor, chain armour, ring armor, ring armour",
+  "491": "chain saw, chainsaw",
+  "492": "chest",
+  "493": "chiffonier, commode",
+  "494": "chime, bell, gong",
+  "495": "china cabinet, china closet",
+  "496": "Christmas stocking",
+  "497": "church, church building",
+  "498": "cinema, movie theater, movie theatre, movie house, picture palace",
+  "499": "cleaver, meat cleaver, chopper",
+  "500": "cliff dwelling",
+  "501": "cloak",
+  "502": "clog, geta, patten, sabot",
+  "503": "cocktail shaker",
+  "504": "coffee mug",
+  "505": "coffeepot",
+  "506": "coil, spiral, volute, whorl, helix",
+  "507": "combination lock",
+  "508": "computer keyboard, keypad",
+  "509": "confectionery, confectionary, candy store",
+  "510": "container ship, containership, container vessel",
+  "511": "convertible",
+  "512": "corkscrew, bottle screw",
+  "513": "cornet, horn, trumpet, trump",
+  "514": "cowboy boot",
+  "515": "cowboy hat, ten-gallon hat",
+  "516": "cradle",
+  "517": "crane",
+  "518": "crash helmet",
+  "519": "crate",
+  "520": "crib, cot",
+  "521": "Crock Pot",
+  "522": "croquet ball",
+  "523": "crutch",
+  "524": "cuirass",
+  "525": "dam, dike, dyke",
+  "526": "desk",
+  "527": "desktop computer",
+  "528": "dial telephone, dial phone",
+  "529": "diaper, nappy, napkin",
+  "530": "digital clock",
+  "531": "digital watch",
+  "532": "dining table, board",
+  "533": "dishrag, dishcloth",
+  "534": "dishwasher, dish washer, dishwashing machine",
+  "535": "disk brake, disc brake",
+  "536": "dock, dockage, docking facility",
+  "537": "dogsled, dog sled, dog sleigh",
+  "538": "dome",
+  "539": "doormat, welcome mat",
+  "540": "drilling platform, offshore rig",
+  "541": "drum, membranophone, tympan",
+  "542": "drumstick",
+  "543": "dumbbell",
+  "544": "Dutch oven",
+  "545": "electric fan, blower",
+  "546": "electric guitar",
+  "547": "electric locomotive",
+  "548": "entertainment center",
+  "549": "envelope",
+  "550": "espresso maker",
+  "551": "face powder",
+  "552": "feather boa, boa",
+  "553": "file, file cabinet, filing cabinet",
+  "554": "fireboat",
+  "555": "fire engine, fire truck",
+  "556": "fire screen, fireguard",
+  "557": "flagpole, flagstaff",
+  "558": "flute, transverse flute",
+  "559": "folding chair",
+  "560": "football helmet",
+  "561": "forklift",
+  "562": "fountain",
+  "563": "fountain pen",
+  "564": "four-poster",
+  "565": "freight car",
+  "566": "French horn, horn",
+  "567": "frying pan, frypan, skillet",
+  "568": "fur coat",
+  "569": "garbage truck, dustcart",
+  "570": "gasmask, respirator, gas helmet",
+  "571": "gas pump, gasoline pump, petrol pump, island dispenser",
+  "572": "goblet",
+  "573": "go-kart",
+  "574": "golf ball",
+  "575": "golfcart, golf cart",
+  "576": "gondola",
+  "577": "gong, tam-tam",
+  "578": "gown",
+  "579": "grand piano, grand",
+  "580": "greenhouse, nursery, glasshouse",
+  "581": "grille, radiator grille",
+  "582": "grocery store, grocery, food market, market",
+  "583": "guillotine",
+  "584": "hair slide",
+  "585": "hair spray",
+  "586": "half track",
+  "587": "hammer",
+  "588": "hamper",
+  "589": "hand blower, blow dryer, blow drier, hair dryer, hair drier",
+  "590": "hand-held computer, hand-held microcomputer",
+  "591": "handkerchief, hankie, hanky, hankey",
+  "592": "hard disc, hard disk, fixed disk",
+  "593": "harmonica, mouth organ, harp, mouth harp",
+  "594": "harp",
+  "595": "harvester, reaper",
+  "596": "hatchet",
+  "597": "holster",
+  "598": "home theater, home theatre",
+  "599": "honeycomb",
+  "600": "hook, claw",
+  "601": "hoopskirt, crinoline",
+  "602": "horizontal bar, high bar",
+  "603": "horse cart, horse-cart",
+  "604": "hourglass",
+  "605": "iPod",
+  "606": "iron, smoothing iron",
+  "607": "jack-o'-lantern",
+  "608": "jean, blue jean, denim",
+  "609": "jeep, landrover",
+  "610": "jersey, T-shirt, tee shirt",
+  "611": "jigsaw puzzle",
+  "612": "jinrikisha, ricksha, rickshaw",
+  "613": "joystick",
+  "614": "kimono",
+  "615": "knee pad",
+  "616": "knot",
+  "617": "lab coat, laboratory coat",
+  "618": "ladle",
+  "619": "lampshade, lamp shade",
+  "620": "laptop, laptop computer",
+  "621": "lawn mower, mower",
+  "622": "lens cap, lens cover",
+  "623": "letter opener, paper knife, paperknife",
+  "624": "library",
+  "625": "lifeboat",
+  "626": "lighter, light, igniter, ignitor",
+  "627": "limousine, limo",
+  "628": "liner, ocean liner",
+  "629": "lipstick, lip rouge",
+  "630": "Loafer",
+  "631": "lotion",
+  "632": "loudspeaker, speaker, speaker unit, loudspeaker system, speaker system",
+  "633": "loupe, jeweler's loupe",
+  "634": "lumbermill, sawmill",
+  "635": "magnetic compass",
+  "636": "mailbag, postbag",
+  "637": "mailbox, letter box",
+  "638": "maillot",
+  "639": "maillot, tank suit",
+  "640": "manhole cover",
+  "641": "maraca",
+  "642": "marimba, xylophone",
+  "643": "mask",
+  "644": "matchstick",
+  "645": "maypole",
+  "646": "maze, labyrinth",
+  "647": "measuring cup",
+  "648": "medicine chest, medicine cabinet",
+  "649": "megalith, megalithic structure",
+  "650": "microphone, mike",
+  "651": "microwave, microwave oven",
+  "652": "military uniform",
+  "653": "milk can",
+  "654": "minibus",
+  "655": "miniskirt, mini",
+  "656": "minivan",
+  "657": "missile",
+  "658": "mitten",
+  "659": "mixing bowl",
+  "660": "mobile home, manufactured home",
+  "661": "Model T",
+  "662": "modem",
+  "663": "monastery",
+  "664": "monitor",
+  "665": "moped",
+  "666": "mortar",
+  "667": "mortarboard",
+  "668": "mosque",
+  "669": "mosquito net",
+  "670": "motor scooter, scooter",
+  "671": "mountain bike, all-terrain bike, off-roader",
+  "672": "mountain tent",
+  "673": "mouse, computer mouse",
+  "674": "mousetrap",
+  "675": "moving van",
+  "676": "muzzle",
+  "677": "nail",
+  "678": "neck brace",
+  "679": "necklace",
+  "680": "nipple",
+  "681": "notebook, notebook computer",
+  "682": "obelisk",
+  "683": "oboe, hautboy, hautbois",
+  "684": "ocarina, sweet potato",
+  "685": "odometer, hodometer, mileometer, milometer",
+  "686": "oil filter",
+  "687": "organ, pipe organ",
+  "688": "oscilloscope, scope, cathode-ray oscilloscope, CRO",
+  "689": "overskirt",
+  "690": "oxcart",
+  "691": "oxygen mask",
+  "692": "packet",
+  "693": "paddle, boat paddle",
+  "694": "paddlewheel, paddle wheel",
+  "695": "padlock",
+  "696": "paintbrush",
+  "697": "pajama, pyjama, pj's, jammies",
+  "698": "palace",
+  "699": "panpipe, pandean pipe, syrinx",
+  "700": "paper towel",
+  "701": "parachute, chute",
+  "702": "parallel bars, bars",
+  "703": "park bench",
+  "704": "parking meter",
+  "705": "passenger car, coach, carriage",
+  "706": "patio, terrace",
+  "707": "pay-phone, pay-station",
+  "708": "pedestal, plinth, footstall",
+  "709": "pencil box, pencil case",
+  "710": "pencil sharpener",
+  "711": "perfume, essence",
+  "712": "Petri dish",
+  "713": "photocopier",
+  "714": "pick, plectrum, plectron",
+  "715": "pickelhaube",
+  "716": "picket fence, paling",
+  "717": "pickup, pickup truck",
+  "718": "pier",
+  "719": "piggy bank, penny bank",
+  "720": "pill bottle",
+  "721": "pillow",
+  "722": "ping-pong ball",
+  "723": "pinwheel",
+  "724": "pirate, pirate ship",
+  "725": "pitcher, ewer",
+  "726": "plane, carpenter's plane, woodworking plane",
+  "727": "planetarium",
+  "728": "plastic bag",
+  "729": "plate rack",
+  "730": "plow, plough",
+  "731": "plunger, plumber's helper",
+  "732": "Polaroid camera, Polaroid Land camera",
+  "733": "pole",
+  "734": "police van, police wagon, paddy wagon, patrol wagon, wagon, black Maria",
+  "735": "poncho",
+  "736": "pool table, billiard table, snooker table",
+  "737": "pop bottle, soda bottle",
+  "738": "pot, flowerpot",
+  "739": "potter's wheel",
+  "740": "power drill",
+  "741": "prayer rug, prayer mat",
+  "742": "printer",
+  "743": "prison, prison house",
+  "744": "projectile, missile",
+  "745": "projector",
+  "746": "puck, hockey puck",
+  "747": "punching bag, punch bag, punching ball, punchball",
+  "748": "purse",
+  "749": "quill, quill pen",
+  "750": "quilt, comforter, comfort, puff",
+  "751": "racer, race car, racing car",
+  "752": "racket, racquet",
+  "753": "radiator",
+  "754": "radio, wireless",
+  "755": "radio telescope, radio reflector",
+  "756": "rain barrel",
+  "757": "recreational vehicle, RV, R.V.",
+  "758": "reel",
+  "759": "reflex camera",
+  "760": "refrigerator, icebox",
+  "761": "remote control, remote",
+  "762": "restaurant, eating house, eating place, eatery",
+  "763": "revolver, six-gun, six-shooter",
+  "764": "rifle",
+  "765": "rocking chair, rocker",
+  "766": "rotisserie",
+  "767": "rubber eraser, rubber, pencil eraser",
+  "768": "rugby ball",
+  "769": "rule, ruler",
+  "770": "running shoe",
+  "771": "safe",
+  "772": "safety pin",
+  "773": "saltshaker, salt shaker",
+  "774": "sandal",
+  "775": "sarong",
+  "776": "sax, saxophone",
+  "777": "scabbard",
+  "778": "scale, weighing machine",
+  "779": "school bus",
+  "780": "schooner",
+  "781": "scoreboard",
+  "782": "screen, CRT screen",
+  "783": "screw",
+  "784": "screwdriver",
+  "785": "seat belt, seatbelt",
+  "786": "sewing machine",
+  "787": "shield, buckler",
+  "788": "shoe shop, shoe-shop, shoe store",
+  "789": "shoji",
+  "790": "shopping basket",
+  "791": "shopping cart",
+  "792": "shovel",
+  "793": "shower cap",
+  "794": "shower curtain",
+  "795": "ski",
+  "796": "ski mask",
+  "797": "sleeping bag",
+  "798": "slide rule, slipstick",
+  "799": "sliding door",
+  "800": "slot, one-armed bandit",
+  "801": "snorkel",
+  "802": "snowmobile",
+  "803": "snowplow, snowplough",
+  "804": "soap dispenser",
+  "805": "soccer ball",
+  "806": "sock",
+  "807": "solar dish, solar collector, solar furnace",
+  "808": "sombrero",
+  "809": "soup bowl",
+  "810": "space bar",
+  "811": "space heater",
+  "812": "space shuttle",
+  "813": "spatula",
+  "814": "speedboat",
+  "815": "spider web, spider's web",
+  "816": "spindle",
+  "817": "sports car, sport car",
+  "818": "spotlight, spot",
+  "819": "stage",
+  "820": "steam locomotive",
+  "821": "steel arch bridge",
+  "822": "steel drum",
+  "823": "stethoscope",
+  "824": "stole",
+  "825": "stone wall",
+  "826": "stopwatch, stop watch",
+  "827": "stove",
+  "828": "strainer",
+  "829": "streetcar, tram, tramcar, trolley, trolley car",
+  "830": "stretcher",
+  "831": "studio couch, day bed",
+  "832": "stupa, tope",
+  "833": "submarine, pigboat, sub, U-boat",
+  "834": "suit, suit of clothes",
+  "835": "sundial",
+  "836": "sunglass",
+  "837": "sunglasses, dark glasses, shades",
+  "838": "sunscreen, sunblock, sun blocker",
+  "839": "suspension bridge",
+  "840": "swab, swob, mop",
+  "841": "sweatshirt",
+  "842": "swimming trunks, bathing trunks",
+  "843": "swing",
+  "844": "switch, electric switch, electrical switch",
+  "845": "syringe",
+  "846": "table lamp",
+  "847": "tank, army tank, armored combat vehicle, armoured combat vehicle",
+  "848": "tape player",
+  "849": "teapot",
+  "850": "teddy, teddy bear",
+  "851": "television, television system",
+  "852": "tennis ball",
+  "853": "thatch, thatched roof",
+  "854": "theater curtain, theatre curtain",
+  "855": "thimble",
+  "856": "thresher, thrasher, threshing machine",
+  "857": "throne",
+  "858": "tile roof",
+  "859": "toaster",
+  "860": "tobacco shop, tobacconist shop, tobacconist",
+  "861": "toilet seat",
+  "862": "torch",
+  "863": "totem pole",
+  "864": "tow truck, tow car, wrecker",
+  "865": "toyshop",
+  "866": "tractor",
+  "867": "trailer truck, tractor trailer, trucking rig, rig, articulated lorry, semi",
+  "868": "tray",
+  "869": "trench coat",
+  "870": "tricycle, trike, velocipede",
+  "871": "trimaran",
+  "872": "tripod",
+  "873": "triumphal arch",
+  "874": "trolleybus, trolley coach, trackless trolley",
+  "875": "trombone",
+  "876": "tub, vat",
+  "877": "turnstile",
+  "878": "typewriter keyboard",
+  "879": "umbrella",
+  "880": "unicycle, monocycle",
+  "881": "upright, upright piano",
+  "882": "vacuum, vacuum cleaner",
+  "883": "vase",
+  "884": "vault",
+  "885": "velvet",
+  "886": "vending machine",
+  "887": "vestment",
+  "888": "viaduct",
+  "889": "violin, fiddle",
+  "890": "volleyball",
+  "891": "waffle iron",
+  "892": "wall clock",
+  "893": "wallet, billfold, notecase, pocketbook",
+  "894": "wardrobe, closet, press",
+  "895": "warplane, military plane",
+  "896": "washbasin, handbasin, washbowl, lavabo, wash-hand basin",
+  "897": "washer, automatic washer, washing machine",
+  "898": "water bottle",
+  "899": "water jug",
+  "900": "water tower",
+  "901": "whiskey jug",
+  "902": "whistle",
+  "903": "wig",
+  "904": "window screen",
+  "905": "window shade",
+  "906": "Windsor tie",
+  "907": "wine bottle",
+  "908": "wing",
+  "909": "wok",
+  "910": "wooden spoon",
+  "911": "wool, woolen, woollen",
+  "912": "worm fence, snake fence, snake-rail fence, Virginia fence",
+  "913": "wreck",
+  "914": "yawl",
+  "915": "yurt",
+  "916": "web site, website, internet site, site",
+  "917": "comic book",
+  "918": "crossword puzzle, crossword",
+  "919": "street sign",
+  "920": "traffic light, traffic signal, stoplight",
+  "921": "book jacket, dust cover, dust jacket, dust wrapper",
+  "922": "menu",
+  "923": "plate",
+  "924": "guacamole",
+  "925": "consomme",
+  "926": "hot pot, hotpot",
+  "927": "trifle",
+  "928": "ice cream, icecream",
+  "929": "ice lolly, lolly, lollipop, popsicle",
+  "930": "French loaf",
+  "931": "bagel, beigel",
+  "932": "pretzel",
+  "933": "cheeseburger",
+  "934": "hotdog, hot dog, red hot",
+  "935": "mashed potato",
+  "936": "head cabbage",
+  "937": "broccoli",
+  "938": "cauliflower",
+  "939": "zucchini, courgette",
+  "940": "spaghetti squash",
+  "941": "acorn squash",
+  "942": "butternut squash",
+  "943": "cucumber, cuke",
+  "944": "artichoke, globe artichoke",
+  "945": "bell pepper",
+  "946": "cardoon",
+  "947": "mushroom",
+  "948": "Granny Smith",
+  "949": "strawberry",
+  "950": "orange",
+  "951": "lemon",
+  "952": "fig",
+  "953": "pineapple, ananas",
+  "954": "banana",
+  "955": "jackfruit, jak, jack",
+  "956": "custard apple",
+  "957": "pomegranate",
+  "958": "hay",
+  "959": "carbonara",
+  "960": "chocolate sauce, chocolate syrup",
+  "961": "dough",
+  "962": "meat loaf, meatloaf",
+  "963": "pizza, pizza pie",
+  "964": "potpie",
+  "965": "burrito",
+  "966": "red wine",
+  "967": "espresso",
+  "968": "cup",
+  "969": "eggnog",
+  "970": "alp",
+  "971": "bubble",
+  "972": "cliff, drop, drop-off",
+  "973": "coral reef",
+  "974": "geyser",
+  "975": "lakeside, lakeshore",
+  "976": "promontory, headland, head, foreland",
+  "977": "sandbar, sand bar",
+  "978": "seashore, coast, seacoast, sea-coast",
+  "979": "valley, vale",
+  "980": "volcano",
+  "981": "ballplayer, baseball player",
+  "982": "groom, bridegroom",
+  "983": "scuba diver",
+  "984": "rapeseed",
+  "985": "daisy",
+  "986": "yellow lady's slipper, yellow lady-slipper, Cypripedium calceolus, Cypripedium parviflorum",
+  "987": "corn",
+  "988": "acorn",
+  "989": "hip, rose hip, rosehip",
+  "990": "buckeye, horse chestnut, conker",
+  "991": "coral fungus",
+  "992": "agaric",
+  "993": "gyromitra",
+  "994": "stinkhorn, carrion fungus",
+  "995": "earthstar",
+  "996": "hen-of-the-woods, hen of the woods, Polyporus frondosus, Grifola frondosa",
+  "997": "bolete",
+  "998": "ear, spike, capitulum",
+  "999": "toilet tissue, toilet paper, bathroom tissue"
+}

requirements.txt

@@ -0,0 +1,5 @@
+torchvision
+safetensors
+pyyaml
+numpy
+pillow

timm/__init__.py

@@ -0,0 +1,4 @@
+from .version import __version__
+from .layers import is_scriptable, is_exportable, set_scriptable, set_exportable
+from .models import create_model, list_models, list_pretrained, is_model, list_modules, model_entrypoint, \
+    is_model_pretrained, get_pretrained_cfg, get_pretrained_cfg_value

timm/data/__init__.py

@@ -0,0 +1,23 @@
+# Copyright 2026 Kiel University
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+# Based on pytorch-image-models (timm); see NOTICE.
+#
+# Modifications:
+# - Pruned exports to the vendored timm subset used by ProgResViT.
+
+from .auto_augment import RandAugment, AutoAugment, rand_augment_ops, auto_augment_policy,\
+    rand_augment_transform, auto_augment_transform
+from .config import resolve_data_config, resolve_model_data_config
+from .constants import *
+from .dataset import ImageDataset, IterableImageDataset, AugMixDataset
+from .dataset_factory import create_dataset
+from .dataset_info import DatasetInfo, CustomDatasetInfo
+from .loader import create_loader
+from .mixup import Mixup, FastCollateMixup
+from .readers import create_reader
+from .readers import get_img_extensions, is_img_extension, set_img_extensions, add_img_extensions, del_img_extensions
+from .real_labels import RealLabelsImagenet
+from .transforms import *
+from .transforms_factory import create_transform

timm/data/auto_augment.py

@@ -0,0 +1,997 @@
+""" AutoAugment, RandAugment, AugMix, and 3-Augment for PyTorch
+
+This code implements the searched ImageNet policies with various tweaks and improvements and
+does not include any of the search code.
+
+AA and RA Implementation adapted from:
+    https://github.com/tensorflow/tpu/blob/master/models/official/efficientnet/autoaugment.py
+
+AugMix adapted from:
+    https://github.com/google-research/augmix
+
+3-Augment based on: https://github.com/facebookresearch/deit/blob/main/README_revenge.md
+
+Papers:
+    AutoAugment: Learning Augmentation Policies from Data - https://arxiv.org/abs/1805.09501
+    Learning Data Augmentation Strategies for Object Detection - https://arxiv.org/abs/1906.11172
+    RandAugment: Practical automated data augmentation... - https://arxiv.org/abs/1909.13719
+    AugMix: A Simple Data Processing Method to Improve Robustness and Uncertainty - https://arxiv.org/abs/1912.02781
+    3-Augment: DeiT III: Revenge of the ViT - https://arxiv.org/abs/2204.07118
+
+Hacked together by / Copyright 2019, Ross Wightman
+"""
+import random
+import math
+import re
+from functools import partial
+from typing import Dict, List, Optional, Union
+
+from PIL import Image, ImageOps, ImageEnhance, ImageChops, ImageFilter
+import PIL
+import numpy as np
+
+
+_PIL_VER = tuple([int(x) for x in PIL.__version__.split('.')[:2]])
+
+_FILL = (128, 128, 128)
+
+_LEVEL_DENOM = 10.  # denominator for conversion from 'Mx' magnitude scale to fractional aug level for op arguments
+
+_HPARAMS_DEFAULT = dict(
+    translate_const=250,
+    img_mean=_FILL,
+)
+
+if hasattr(Image, "Resampling"):
+    _RANDOM_INTERPOLATION = (Image.Resampling.BILINEAR, Image.Resampling.BICUBIC)
+    _DEFAULT_INTERPOLATION = Image.Resampling.BICUBIC
+else:
+    _RANDOM_INTERPOLATION = (Image.BILINEAR, Image.BICUBIC)
+    _DEFAULT_INTERPOLATION = Image.BICUBIC
+
+
+def _interpolation(kwargs):
+    interpolation = kwargs.pop('resample', _DEFAULT_INTERPOLATION)
+    if isinstance(interpolation, (list, tuple)):
+        return random.choice(interpolation)
+    return interpolation
+
+
+def _check_args_tf(kwargs):
+    if 'fillcolor' in kwargs and _PIL_VER < (5, 0):
+        kwargs.pop('fillcolor')
+    kwargs['resample'] = _interpolation(kwargs)
+
+
+def shear_x(img, factor, **kwargs):
+    _check_args_tf(kwargs)
+    return img.transform(img.size, Image.AFFINE, (1, factor, 0, 0, 1, 0), **kwargs)
+
+
+def shear_y(img, factor, **kwargs):
+    _check_args_tf(kwargs)
+    return img.transform(img.size, Image.AFFINE, (1, 0, 0, factor, 1, 0), **kwargs)
+
+
+def translate_x_rel(img, pct, **kwargs):
+    pixels = pct * img.size[0]
+    _check_args_tf(kwargs)
+    return img.transform(img.size, Image.AFFINE, (1, 0, pixels, 0, 1, 0), **kwargs)
+
+
+def translate_y_rel(img, pct, **kwargs):
+    pixels = pct * img.size[1]
+    _check_args_tf(kwargs)
+    return img.transform(img.size, Image.AFFINE, (1, 0, 0, 0, 1, pixels), **kwargs)
+
+
+def translate_x_abs(img, pixels, **kwargs):
+    _check_args_tf(kwargs)
+    return img.transform(img.size, Image.AFFINE, (1, 0, pixels, 0, 1, 0), **kwargs)
+
+
+def translate_y_abs(img, pixels, **kwargs):
+    _check_args_tf(kwargs)
+    return img.transform(img.size, Image.AFFINE, (1, 0, 0, 0, 1, pixels), **kwargs)
+
+
+def rotate(img, degrees, **kwargs):
+    _check_args_tf(kwargs)
+    if _PIL_VER >= (5, 2):
+        return img.rotate(degrees, **kwargs)
+    if _PIL_VER >= (5, 0):
+        w, h = img.size
+        post_trans = (0, 0)
+        rotn_center = (w / 2.0, h / 2.0)
+        angle = -math.radians(degrees)
+        matrix = [
+            round(math.cos(angle), 15),
+            round(math.sin(angle), 15),
+            0.0,
+            round(-math.sin(angle), 15),
+            round(math.cos(angle), 15),
+            0.0,
+        ]
+
+        def transform(x, y, matrix):
+            (a, b, c, d, e, f) = matrix
+            return a * x + b * y + c, d * x + e * y + f
+
+        matrix[2], matrix[5] = transform(
+            -rotn_center[0] - post_trans[0], -rotn_center[1] - post_trans[1], matrix
+        )
+        matrix[2] += rotn_center[0]
+        matrix[5] += rotn_center[1]
+        return img.transform(img.size, Image.AFFINE, matrix, **kwargs)
+    return img.rotate(degrees, resample=kwargs['resample'])
+
+
+def auto_contrast(img, **__):
+    return ImageOps.autocontrast(img)
+
+
+def invert(img, **__):
+    return ImageOps.invert(img)
+
+
+def equalize(img, **__):
+    return ImageOps.equalize(img)
+
+
+def solarize(img, thresh, **__):
+    return ImageOps.solarize(img, thresh)
+
+
+def solarize_add(img, add, thresh=128, **__):
+    lut = []
+    for i in range(256):
+        if i < thresh:
+            lut.append(min(255, i + add))
+        else:
+            lut.append(i)
+
+    if img.mode in ("L", "RGB"):
+        if img.mode == "RGB" and len(lut) == 256:
+            lut = lut + lut + lut
+        return img.point(lut)
+
+    return img
+
+
+def posterize(img, bits_to_keep, **__):
+    if bits_to_keep >= 8:
+        return img
+    return ImageOps.posterize(img, bits_to_keep)
+
+
+def contrast(img, factor, **__):
+    return ImageEnhance.Contrast(img).enhance(factor)
+
+
+def color(img, factor, **__):
+    return ImageEnhance.Color(img).enhance(factor)
+
+
+def brightness(img, factor, **__):
+    return ImageEnhance.Brightness(img).enhance(factor)
+
+
+def sharpness(img, factor, **__):
+    return ImageEnhance.Sharpness(img).enhance(factor)
+
+
+def gaussian_blur(img, factor, **__):
+    img = img.filter(ImageFilter.GaussianBlur(radius=factor))
+    return img
+
+
+def gaussian_blur_rand(img, factor, **__):
+    radius_min = 0.1
+    radius_max = 2.0
+    img = img.filter(ImageFilter.GaussianBlur(radius=random.uniform(radius_min, radius_max * factor)))
+    return img
+
+
+def desaturate(img, factor, **_):
+    factor = min(1., max(0., 1. - factor))
+    # enhance factor 0 = grayscale, 1.0 = no-change
+    return ImageEnhance.Color(img).enhance(factor)
+
+
+def _randomly_negate(v):
+    """With 50% prob, negate the value"""
+    return -v if random.random() > 0.5 else v
+
+
+def _rotate_level_to_arg(level, _hparams):
+    # range [-30, 30]
+    level = (level / _LEVEL_DENOM) * 30.
+    level = _randomly_negate(level)
+    return level,
+
+
+def _enhance_level_to_arg(level, _hparams):
+    # range [0.1, 1.9]
+    return (level / _LEVEL_DENOM) * 1.8 + 0.1,
+
+
+def _enhance_increasing_level_to_arg(level, _hparams):
+    # the 'no change' level is 1.0, moving away from that towards 0. or 2.0 increases the enhancement blend
+    # range [0.1, 1.9] if level <= _LEVEL_DENOM
+    level = (level / _LEVEL_DENOM) * .9
+    level = max(0.1, 1.0 + _randomly_negate(level))  # keep it >= 0.1
+    return level,
+
+
+def _minmax_level_to_arg(level, _hparams, min_val=0., max_val=1.0, clamp=True):
+    level = (level / _LEVEL_DENOM)
+    level = min_val + (max_val - min_val) * level
+    if clamp:
+        level = max(min_val, min(max_val, level))
+    return level,
+
+
+def _shear_level_to_arg(level, _hparams):
+    # range [-0.3, 0.3]
+    level = (level / _LEVEL_DENOM) * 0.3
+    level = _randomly_negate(level)
+    return level,
+
+
+def _translate_abs_level_to_arg(level, hparams):
+    translate_const = hparams['translate_const']
+    level = (level / _LEVEL_DENOM) * float(translate_const)
+    level = _randomly_negate(level)
+    return level,
+
+
+def _translate_rel_level_to_arg(level, hparams):
+    # default range [-0.45, 0.45]
+    translate_pct = hparams.get('translate_pct', 0.45)
+    level = (level / _LEVEL_DENOM) * translate_pct
+    level = _randomly_negate(level)
+    return level,
+
+
+def _posterize_level_to_arg(level, _hparams):
+    # As per Tensorflow TPU EfficientNet impl
+    # range [0, 4], 'keep 0 up to 4 MSB of original image'
+    # intensity/severity of augmentation decreases with level
+    return int((level / _LEVEL_DENOM) * 4),
+
+
+def _posterize_increasing_level_to_arg(level, hparams):
+    # As per Tensorflow models research and UDA impl
+    # range [4, 0], 'keep 4 down to 0 MSB of original image',
+    # intensity/severity of augmentation increases with level
+    return 4 - _posterize_level_to_arg(level, hparams)[0],
+
+
+def _posterize_original_level_to_arg(level, _hparams):
+    # As per original AutoAugment paper description
+    # range [4, 8], 'keep 4 up to 8 MSB of image'
+    # intensity/severity of augmentation decreases with level
+    return int((level / _LEVEL_DENOM) * 4) + 4,
+
+
+def _solarize_level_to_arg(level, _hparams):
+    # range [0, 256]
+    # intensity/severity of augmentation decreases with level
+    return min(256, int((level / _LEVEL_DENOM) * 256)),
+
+
+def _solarize_increasing_level_to_arg(level, _hparams):
+    # range [0, 256]
+    # intensity/severity of augmentation increases with level
+    return 256 - _solarize_level_to_arg(level, _hparams)[0],
+
+
+def _solarize_add_level_to_arg(level, _hparams):
+    # range [0, 110]
+    return min(128, int((level / _LEVEL_DENOM) * 110)),
+
+
+LEVEL_TO_ARG = {
+    'AutoContrast': None,
+    'Equalize': None,
+    'Invert': None,
+    'Rotate': _rotate_level_to_arg,
+    # There are several variations of the posterize level scaling in various Tensorflow/Google repositories/papers
+    'Posterize': _posterize_level_to_arg,
+    'PosterizeIncreasing': _posterize_increasing_level_to_arg,
+    'PosterizeOriginal': _posterize_original_level_to_arg,
+    'Solarize': _solarize_level_to_arg,
+    'SolarizeIncreasing': _solarize_increasing_level_to_arg,
+    'SolarizeAdd': _solarize_add_level_to_arg,
+    'Color': _enhance_level_to_arg,
+    'ColorIncreasing': _enhance_increasing_level_to_arg,
+    'Contrast': _enhance_level_to_arg,
+    'ContrastIncreasing': _enhance_increasing_level_to_arg,
+    'Brightness': _enhance_level_to_arg,
+    'BrightnessIncreasing': _enhance_increasing_level_to_arg,
+    'Sharpness': _enhance_level_to_arg,
+    'SharpnessIncreasing': _enhance_increasing_level_to_arg,
+    'ShearX': _shear_level_to_arg,
+    'ShearY': _shear_level_to_arg,
+    'TranslateX': _translate_abs_level_to_arg,
+    'TranslateY': _translate_abs_level_to_arg,
+    'TranslateXRel': _translate_rel_level_to_arg,
+    'TranslateYRel': _translate_rel_level_to_arg,
+    'Desaturate': partial(_minmax_level_to_arg, min_val=0.5, max_val=1.0),
+    'GaussianBlur': partial(_minmax_level_to_arg, min_val=0.1, max_val=2.0),
+    'GaussianBlurRand': _minmax_level_to_arg,
+}
+
+
+NAME_TO_OP = {
+    'AutoContrast': auto_contrast,
+    'Equalize': equalize,
+    'Invert': invert,
+    'Rotate': rotate,
+    'Posterize': posterize,
+    'PosterizeIncreasing': posterize,
+    'PosterizeOriginal': posterize,
+    'Solarize': solarize,
+    'SolarizeIncreasing': solarize,
+    'SolarizeAdd': solarize_add,
+    'Color': color,
+    'ColorIncreasing': color,
+    'Contrast': contrast,
+    'ContrastIncreasing': contrast,
+    'Brightness': brightness,
+    'BrightnessIncreasing': brightness,
+    'Sharpness': sharpness,
+    'SharpnessIncreasing': sharpness,
+    'ShearX': shear_x,
+    'ShearY': shear_y,
+    'TranslateX': translate_x_abs,
+    'TranslateY': translate_y_abs,
+    'TranslateXRel': translate_x_rel,
+    'TranslateYRel': translate_y_rel,
+    'Desaturate': desaturate,
+    'GaussianBlur': gaussian_blur,
+    'GaussianBlurRand': gaussian_blur_rand,
+}
+
+
+class AugmentOp:
+
+    def __init__(self, name, prob=0.5, magnitude=10, hparams=None):
+        hparams = hparams or _HPARAMS_DEFAULT
+        self.name = name
+        self.aug_fn = NAME_TO_OP[name]
+        self.level_fn = LEVEL_TO_ARG[name]
+        self.prob = prob
+        self.magnitude = magnitude
+        self.hparams = hparams.copy()
+        self.kwargs = dict(
+            fillcolor=hparams['img_mean'] if 'img_mean' in hparams else _FILL,
+            resample=hparams['interpolation'] if 'interpolation' in hparams else _RANDOM_INTERPOLATION,
+        )
+
+        # If magnitude_std is > 0, we introduce some randomness
+        # in the usually fixed policy and sample magnitude from a normal distribution
+        # with mean `magnitude` and std-dev of `magnitude_std`.
+        # NOTE This is my own hack, being tested, not in papers or reference impls.
+        # If magnitude_std is inf, we sample magnitude from a uniform distribution
+        self.magnitude_std = self.hparams.get('magnitude_std', 0)
+        self.magnitude_max = self.hparams.get('magnitude_max', None)
+
+    def __call__(self, img):
+        if self.prob < 1.0 and random.random() > self.prob:
+            return img
+        magnitude = self.magnitude
+        if self.magnitude_std > 0:
+            # magnitude randomization enabled
+            if self.magnitude_std == float('inf'):
+                # inf == uniform sampling
+                magnitude = random.uniform(0, magnitude)
+            elif self.magnitude_std > 0:
+                magnitude = random.gauss(magnitude, self.magnitude_std)
+        # default upper_bound for the timm RA impl is _LEVEL_DENOM (10)
+        # setting magnitude_max overrides this to allow M > 10 (behaviour closer to Google TF RA impl)
+        upper_bound = self.magnitude_max or _LEVEL_DENOM
+        magnitude = max(0., min(magnitude, upper_bound))
+        level_args = self.level_fn(magnitude, self.hparams) if self.level_fn is not None else tuple()
+        return self.aug_fn(img, *level_args, **self.kwargs)
+
+    def __repr__(self):
+        fs = self.__class__.__name__ + f'(name={self.name}, p={self.prob}'
+        fs += f', m={self.magnitude}, mstd={self.magnitude_std}'
+        if self.magnitude_max is not None:
+            fs += f', mmax={self.magnitude_max}'
+        fs += ')'
+        return fs
+
+
+def auto_augment_policy_v0(hparams):
+    # ImageNet v0 policy from TPU EfficientNet impl, cannot find a paper reference.
+    policy = [
+        [('Equalize', 0.8, 1), ('ShearY', 0.8, 4)],
+        [('Color', 0.4, 9), ('Equalize', 0.6, 3)],
+        [('Color', 0.4, 1), ('Rotate', 0.6, 8)],
+        [('Solarize', 0.8, 3), ('Equalize', 0.4, 7)],
+        [('Solarize', 0.4, 2), ('Solarize', 0.6, 2)],
+        [('Color', 0.2, 0), ('Equalize', 0.8, 8)],
+        [('Equalize', 0.4, 8), ('SolarizeAdd', 0.8, 3)],
+        [('ShearX', 0.2, 9), ('Rotate', 0.6, 8)],
+        [('Color', 0.6, 1), ('Equalize', 1.0, 2)],
+        [('Invert', 0.4, 9), ('Rotate', 0.6, 0)],
+        [('Equalize', 1.0, 9), ('ShearY', 0.6, 3)],
+        [('Color', 0.4, 7), ('Equalize', 0.6, 0)],
+        [('Posterize', 0.4, 6), ('AutoContrast', 0.4, 7)],
+        [('Solarize', 0.6, 8), ('Color', 0.6, 9)],
+        [('Solarize', 0.2, 4), ('Rotate', 0.8, 9)],
+        [('Rotate', 1.0, 7), ('TranslateYRel', 0.8, 9)],
+        [('ShearX', 0.0, 0), ('Solarize', 0.8, 4)],
+        [('ShearY', 0.8, 0), ('Color', 0.6, 4)],
+        [('Color', 1.0, 0), ('Rotate', 0.6, 2)],
+        [('Equalize', 0.8, 4), ('Equalize', 0.0, 8)],
+        [('Equalize', 1.0, 4), ('AutoContrast', 0.6, 2)],
+        [('ShearY', 0.4, 7), ('SolarizeAdd', 0.6, 7)],
+        [('Posterize', 0.8, 2), ('Solarize', 0.6, 10)],  # This results in black image with Tpu posterize
+        [('Solarize', 0.6, 8), ('Equalize', 0.6, 1)],
+        [('Color', 0.8, 6), ('Rotate', 0.4, 5)],
+    ]
+    pc = [[AugmentOp(*a, hparams=hparams) for a in sp] for sp in policy]
+    return pc
+
+
+def auto_augment_policy_v0r(hparams):
+    # ImageNet v0 policy from TPU EfficientNet impl, with variation of Posterize used
+    # in Google research implementation (number of bits discarded increases with magnitude)
+    policy = [
+        [('Equalize', 0.8, 1), ('ShearY', 0.8, 4)],
+        [('Color', 0.4, 9), ('Equalize', 0.6, 3)],
+        [('Color', 0.4, 1), ('Rotate', 0.6, 8)],
+        [('Solarize', 0.8, 3), ('Equalize', 0.4, 7)],
+        [('Solarize', 0.4, 2), ('Solarize', 0.6, 2)],
+        [('Color', 0.2, 0), ('Equalize', 0.8, 8)],
+        [('Equalize', 0.4, 8), ('SolarizeAdd', 0.8, 3)],
+        [('ShearX', 0.2, 9), ('Rotate', 0.6, 8)],
+        [('Color', 0.6, 1), ('Equalize', 1.0, 2)],
+        [('Invert', 0.4, 9), ('Rotate', 0.6, 0)],
+        [('Equalize', 1.0, 9), ('ShearY', 0.6, 3)],
+        [('Color', 0.4, 7), ('Equalize', 0.6, 0)],
+        [('PosterizeIncreasing', 0.4, 6), ('AutoContrast', 0.4, 7)],
+        [('Solarize', 0.6, 8), ('Color', 0.6, 9)],
+        [('Solarize', 0.2, 4), ('Rotate', 0.8, 9)],
+        [('Rotate', 1.0, 7), ('TranslateYRel', 0.8, 9)],
+        [('ShearX', 0.0, 0), ('Solarize', 0.8, 4)],
+        [('ShearY', 0.8, 0), ('Color', 0.6, 4)],
+        [('Color', 1.0, 0), ('Rotate', 0.6, 2)],
+        [('Equalize', 0.8, 4), ('Equalize', 0.0, 8)],
+        [('Equalize', 1.0, 4), ('AutoContrast', 0.6, 2)],
+        [('ShearY', 0.4, 7), ('SolarizeAdd', 0.6, 7)],
+        [('PosterizeIncreasing', 0.8, 2), ('Solarize', 0.6, 10)],
+        [('Solarize', 0.6, 8), ('Equalize', 0.6, 1)],
+        [('Color', 0.8, 6), ('Rotate', 0.4, 5)],
+    ]
+    pc = [[AugmentOp(*a, hparams=hparams) for a in sp] for sp in policy]
+    return pc
+
+
+def auto_augment_policy_original(hparams):
+    # ImageNet policy from https://arxiv.org/abs/1805.09501
+    policy = [
+        [('PosterizeOriginal', 0.4, 8), ('Rotate', 0.6, 9)],
+        [('Solarize', 0.6, 5), ('AutoContrast', 0.6, 5)],
+        [('Equalize', 0.8, 8), ('Equalize', 0.6, 3)],
+        [('PosterizeOriginal', 0.6, 7), ('PosterizeOriginal', 0.6, 6)],
+        [('Equalize', 0.4, 7), ('Solarize', 0.2, 4)],
+        [('Equalize', 0.4, 4), ('Rotate', 0.8, 8)],
+        [('Solarize', 0.6, 3), ('Equalize', 0.6, 7)],
+        [('PosterizeOriginal', 0.8, 5), ('Equalize', 1.0, 2)],
+        [('Rotate', 0.2, 3), ('Solarize', 0.6, 8)],
+        [('Equalize', 0.6, 8), ('PosterizeOriginal', 0.4, 6)],
+        [('Rotate', 0.8, 8), ('Color', 0.4, 0)],
+        [('Rotate', 0.4, 9), ('Equalize', 0.6, 2)],
+        [('Equalize', 0.0, 7), ('Equalize', 0.8, 8)],
+        [('Invert', 0.6, 4), ('Equalize', 1.0, 8)],
+        [('Color', 0.6, 4), ('Contrast', 1.0, 8)],
+        [('Rotate', 0.8, 8), ('Color', 1.0, 2)],
+        [('Color', 0.8, 8), ('Solarize', 0.8, 7)],
+        [('Sharpness', 0.4, 7), ('Invert', 0.6, 8)],
+        [('ShearX', 0.6, 5), ('Equalize', 1.0, 9)],
+        [('Color', 0.4, 0), ('Equalize', 0.6, 3)],
+        [('Equalize', 0.4, 7), ('Solarize', 0.2, 4)],
+        [('Solarize', 0.6, 5), ('AutoContrast', 0.6, 5)],
+        [('Invert', 0.6, 4), ('Equalize', 1.0, 8)],
+        [('Color', 0.6, 4), ('Contrast', 1.0, 8)],
+        [('Equalize', 0.8, 8), ('Equalize', 0.6, 3)],
+    ]
+    pc = [[AugmentOp(*a, hparams=hparams) for a in sp] for sp in policy]
+    return pc
+
+
+def auto_augment_policy_originalr(hparams):
+    # ImageNet policy from https://arxiv.org/abs/1805.09501 with research posterize variation
+    policy = [
+        [('PosterizeIncreasing', 0.4, 8), ('Rotate', 0.6, 9)],
+        [('Solarize', 0.6, 5), ('AutoContrast', 0.6, 5)],
+        [('Equalize', 0.8, 8), ('Equalize', 0.6, 3)],
+        [('PosterizeIncreasing', 0.6, 7), ('PosterizeIncreasing', 0.6, 6)],
+        [('Equalize', 0.4, 7), ('Solarize', 0.2, 4)],
+        [('Equalize', 0.4, 4), ('Rotate', 0.8, 8)],
+        [('Solarize', 0.6, 3), ('Equalize', 0.6, 7)],
+        [('PosterizeIncreasing', 0.8, 5), ('Equalize', 1.0, 2)],
+        [('Rotate', 0.2, 3), ('Solarize', 0.6, 8)],
+        [('Equalize', 0.6, 8), ('PosterizeIncreasing', 0.4, 6)],
+        [('Rotate', 0.8, 8), ('Color', 0.4, 0)],
+        [('Rotate', 0.4, 9), ('Equalize', 0.6, 2)],
+        [('Equalize', 0.0, 7), ('Equalize', 0.8, 8)],
+        [('Invert', 0.6, 4), ('Equalize', 1.0, 8)],
+        [('Color', 0.6, 4), ('Contrast', 1.0, 8)],
+        [('Rotate', 0.8, 8), ('Color', 1.0, 2)],
+        [('Color', 0.8, 8), ('Solarize', 0.8, 7)],
+        [('Sharpness', 0.4, 7), ('Invert', 0.6, 8)],
+        [('ShearX', 0.6, 5), ('Equalize', 1.0, 9)],
+        [('Color', 0.4, 0), ('Equalize', 0.6, 3)],
+        [('Equalize', 0.4, 7), ('Solarize', 0.2, 4)],
+        [('Solarize', 0.6, 5), ('AutoContrast', 0.6, 5)],
+        [('Invert', 0.6, 4), ('Equalize', 1.0, 8)],
+        [('Color', 0.6, 4), ('Contrast', 1.0, 8)],
+        [('Equalize', 0.8, 8), ('Equalize', 0.6, 3)],
+    ]
+    pc = [[AugmentOp(*a, hparams=hparams) for a in sp] for sp in policy]
+    return pc
+
+
+def auto_augment_policy_3a(hparams):
+    policy = [
+        [('Solarize', 1.0, 5)],  # 128 solarize threshold @ 5 magnitude
+        [('Desaturate', 1.0, 10)],  # grayscale at 10 magnitude
+        [('GaussianBlurRand', 1.0, 10)],
+    ]
+    pc = [[AugmentOp(*a, hparams=hparams) for a in sp] for sp in policy]
+    return pc
+
+
+def auto_augment_policy(name='v0', hparams=None):
+    hparams = hparams or _HPARAMS_DEFAULT
+    if name == 'original':
+        return auto_augment_policy_original(hparams)
+    if name == 'originalr':
+        return auto_augment_policy_originalr(hparams)
+    if name == 'v0':
+        return auto_augment_policy_v0(hparams)
+    if name == 'v0r':
+        return auto_augment_policy_v0r(hparams)
+    if name == '3a':
+        return auto_augment_policy_3a(hparams)
+    assert False, f'Unknown AA policy {name}'
+
+
+class AutoAugment:
+
+    def __init__(self, policy):
+        self.policy = policy
+
+    def __call__(self, img):
+        sub_policy = random.choice(self.policy)
+        for op in sub_policy:
+            img = op(img)
+        return img
+
+    def __repr__(self):
+        fs = self.__class__.__name__ + '(policy='
+        for p in self.policy:
+            fs += '\n\t['
+            fs += ', '.join([str(op) for op in p])
+            fs += ']'
+        fs += ')'
+        return fs
+
+
+def auto_augment_transform(config_str: str, hparams: Optional[Dict] = None):
+    """
+    Create a AutoAugment transform
+
+    Args:
+        config_str: String defining configuration of auto augmentation. Consists of multiple sections separated by
+            dashes ('-').
+            The first section defines the AutoAugment policy (one of 'v0', 'v0r', 'original', 'originalr').
+
+            The remaining sections:
+                'mstd' -  float std deviation of magnitude noise applied
+            Ex 'original-mstd0.5' results in AutoAugment with original policy, magnitude_std 0.5
+
+        hparams: Other hparams (kwargs) for the AutoAugmentation scheme
+
+    Returns:
+         A PyTorch compatible Transform
+    """
+    config = config_str.split('-')
+    policy_name = config[0]
+    config = config[1:]
+    for c in config:
+        cs = re.split(r'(\d.*)', c)
+        if len(cs) < 2:
+            continue
+        key, val = cs[:2]
+        if key == 'mstd':
+            # noise param injected via hparams for now
+            hparams.setdefault('magnitude_std', float(val))
+        else:
+            assert False, 'Unknown AutoAugment config section'
+    aa_policy = auto_augment_policy(policy_name, hparams=hparams)
+    return AutoAugment(aa_policy)
+
+
+_RAND_TRANSFORMS = [
+    'AutoContrast',
+    'Equalize',
+    'Invert',
+    'Rotate',
+    'Posterize',
+    'Solarize',
+    'SolarizeAdd',
+    'Color',
+    'Contrast',
+    'Brightness',
+    'Sharpness',
+    'ShearX',
+    'ShearY',
+    'TranslateXRel',
+    'TranslateYRel',
+    # 'Cutout'  # NOTE I've implement this as random erasing separately
+]
+
+
+_RAND_INCREASING_TRANSFORMS = [
+    'AutoContrast',
+    'Equalize',
+    'Invert',
+    'Rotate',
+    'PosterizeIncreasing',
+    'SolarizeIncreasing',
+    'SolarizeAdd',
+    'ColorIncreasing',
+    'ContrastIncreasing',
+    'BrightnessIncreasing',
+    'SharpnessIncreasing',
+    'ShearX',
+    'ShearY',
+    'TranslateXRel',
+    'TranslateYRel',
+    # 'Cutout'  # NOTE I've implement this as random erasing separately
+]
+
+
+_RAND_3A = [
+    'SolarizeIncreasing',
+    'Desaturate',
+    'GaussianBlur',
+]
+
+
+_RAND_WEIGHTED_3A = {
+    'SolarizeIncreasing': 6,
+    'Desaturate': 6,
+    'GaussianBlur': 6,
+    'Rotate': 3,
+    'ShearX': 2,
+    'ShearY': 2,
+    'PosterizeIncreasing': 1,
+    'AutoContrast': 1,
+    'ColorIncreasing': 1,
+    'SharpnessIncreasing': 1,
+    'ContrastIncreasing': 1,
+    'BrightnessIncreasing': 1,
+    'Equalize': 1,
+    'Invert': 1,
+}
+
+
+# These experimental weights are based loosely on the relative improvements mentioned in paper.
+# They may not result in increased performance, but could likely be tuned to so.
+_RAND_WEIGHTED_0 = {
+    'Rotate': 3,
+    'ShearX': 2,
+    'ShearY': 2,
+    'TranslateXRel': 1,
+    'TranslateYRel': 1,
+    'ColorIncreasing': .25,
+    'SharpnessIncreasing': 0.25,
+    'AutoContrast': 0.25,
+    'SolarizeIncreasing': .05,
+    'SolarizeAdd': .05,
+    'ContrastIncreasing': .05,
+    'BrightnessIncreasing': .05,
+    'Equalize': .05,
+    'PosterizeIncreasing': 0.05,
+    'Invert': 0.05,
+}
+
+
+def _get_weighted_transforms(transforms: Dict):
+    transforms, probs = list(zip(*transforms.items()))
+    probs = np.array(probs)
+    probs = probs / np.sum(probs)
+    return transforms, probs
+
+
+def rand_augment_choices(name: str, increasing=True):
+    if name == 'weights':
+        return _RAND_WEIGHTED_0
+    if name == '3aw':
+        return _RAND_WEIGHTED_3A
+    if name == '3a':
+        return _RAND_3A
+    return _RAND_INCREASING_TRANSFORMS if increasing else _RAND_TRANSFORMS
+
+
+def rand_augment_ops(
+        magnitude: Union[int, float] = 10,
+        prob: float = 0.5,
+        hparams: Optional[Dict] = None,
+        transforms: Optional[Union[Dict, List]] = None,
+):
+    hparams = hparams or _HPARAMS_DEFAULT
+    transforms = transforms or _RAND_TRANSFORMS
+    return [AugmentOp(
+        name, prob=prob, magnitude=magnitude, hparams=hparams) for name in transforms]
+
+
+class RandAugment:
+    def __init__(self, ops, num_layers=2, choice_weights=None):
+        self.ops = ops
+        self.num_layers = num_layers
+        self.choice_weights = choice_weights
+
+    def __call__(self, img):
+        # no replacement when using weighted choice
+        ops = np.random.choice(
+            self.ops,
+            self.num_layers,
+            replace=self.choice_weights is None,
+            p=self.choice_weights,
+        )
+        for op in ops:
+            img = op(img)
+        return img
+
+    def __repr__(self):
+        fs = self.__class__.__name__ + f'(n={self.num_layers}, ops='
+        for op in self.ops:
+            fs += f'\n\t{op}'
+        fs += ')'
+        return fs
+
+
+def rand_augment_transform(
+        config_str: str,
+        hparams: Optional[Dict] = None,
+        transforms: Optional[Union[str, Dict, List]] = None,
+):
+    """
+    Create a RandAugment transform
+
+    Args:
+        config_str (str): String defining configuration of random augmentation. Consists of multiple sections separated
+            by dashes ('-'). The first section defines the specific variant of rand augment (currently only 'rand').
+            The remaining sections, not order sepecific determine
+                'm' - integer magnitude of rand augment
+                'n' - integer num layers (number of transform ops selected per image)
+                'p' - float probability of applying each layer (default 0.5)
+                'mstd' -  float std deviation of magnitude noise applied, or uniform sampling if infinity (or > 100)
+                'mmax' - set upper bound for magnitude to something other than default of  _LEVEL_DENOM (10)
+                'inc' - integer (bool), use augmentations that increase in severity with magnitude (default: 0)
+                't' - str name of transform set to use
+            Ex 'rand-m9-n3-mstd0.5' results in RandAugment with magnitude 9, num_layers 3, magnitude_std 0.5
+            'rand-mstd1-tweights' results in mag std 1.0, weighted transforms, default mag of 10 and num_layers 2
+
+        hparams (dict): Other hparams (kwargs) for the RandAugmentation scheme
+
+    Returns:
+         A PyTorch compatible Transform
+    """
+    magnitude = _LEVEL_DENOM  # default to _LEVEL_DENOM for magnitude (currently 10)
+    num_layers = 2  # default to 2 ops per image
+    increasing = False
+    prob = 0.5
+    config = config_str.split('-')
+    assert config[0] == 'rand'
+    config = config[1:]
+    for c in config:
+        if c.startswith('t'):
+            # NOTE old 'w' key was removed, 'w0' is not equivalent to 'tweights'
+            val = str(c[1:])
+            if transforms is None:
+                transforms = val
+        else:
+            # numeric options
+            cs = re.split(r'(\d.*)', c)
+            if len(cs) < 2:
+                continue
+            key, val = cs[:2]
+            if key == 'mstd':
+                # noise param / randomization of magnitude values
+                mstd = float(val)
+                if mstd > 100:
+                    # use uniform sampling in 0 to magnitude if mstd is > 100
+                    mstd = float('inf')
+                hparams.setdefault('magnitude_std', mstd)
+            elif key == 'mmax':
+                # clip magnitude between [0, mmax] instead of default [0, _LEVEL_DENOM]
+                hparams.setdefault('magnitude_max', int(val))
+            elif key == 'inc':
+                if bool(val):
+                    increasing = True
+            elif key == 'm':
+                magnitude = int(val)
+            elif key == 'n':
+                num_layers = int(val)
+            elif key == 'p':
+                prob = float(val)
+            else:
+                assert False, 'Unknown RandAugment config section'
+
+    if isinstance(transforms, str):
+        transforms = rand_augment_choices(transforms, increasing=increasing)
+    elif transforms is None:
+        transforms = _RAND_INCREASING_TRANSFORMS if increasing else _RAND_TRANSFORMS
+
+    choice_weights = None
+    if isinstance(transforms, Dict):
+        transforms, choice_weights = _get_weighted_transforms(transforms)
+
+    ra_ops = rand_augment_ops(magnitude=magnitude, prob=prob, hparams=hparams, transforms=transforms)
+    return RandAugment(ra_ops, num_layers, choice_weights=choice_weights)
+
+
+_AUGMIX_TRANSFORMS = [
+    'AutoContrast',
+    'ColorIncreasing',  # not in paper
+    'ContrastIncreasing',  # not in paper
+    'BrightnessIncreasing',  # not in paper
+    'SharpnessIncreasing',  # not in paper
+    'Equalize',
+    'Rotate',
+    'PosterizeIncreasing',
+    'SolarizeIncreasing',
+    'ShearX',
+    'ShearY',
+    'TranslateXRel',
+    'TranslateYRel',
+]
+
+
+def augmix_ops(
+        magnitude: Union[int, float] = 10,
+        hparams: Optional[Dict] = None,
+        transforms: Optional[Union[str, Dict, List]] = None,
+):
+    hparams = hparams or _HPARAMS_DEFAULT
+    transforms = transforms or _AUGMIX_TRANSFORMS
+    return [AugmentOp(
+        name,
+        prob=1.0,
+        magnitude=magnitude,
+        hparams=hparams
+    ) for name in transforms]
+
+
+class AugMixAugment:
+    """ AugMix Transform
+    Adapted and improved from impl here: https://github.com/google-research/augmix/blob/master/imagenet.py
+    From paper: 'AugMix: A Simple Data Processing Method to Improve Robustness and Uncertainty -
+    https://arxiv.org/abs/1912.02781
+    """
+    def __init__(self, ops, alpha=1., width=3, depth=-1, blended=False):
+        self.ops = ops
+        self.alpha = alpha
+        self.width = width
+        self.depth = depth
+        self.blended = blended  # blended mode is faster but not well tested
+
+    def _calc_blended_weights(self, ws, m):
+        ws = ws * m
+        cump = 1.
+        rws = []
+        for w in ws[::-1]:
+            alpha = w / cump
+            cump *= (1 - alpha)
+            rws.append(alpha)
+        return np.array(rws[::-1], dtype=np.float32)
+
+    def _apply_blended(self, img, mixing_weights, m):
+        # This is my first crack and implementing a slightly faster mixed augmentation. Instead
+        # of accumulating the mix for each chain in a Numpy array and then blending with original,
+        # it recomputes the blending coefficients and applies one PIL image blend per chain.
+        # TODO the results appear in the right ballpark but they differ by more than rounding.
+        img_orig = img.copy()
+        ws = self._calc_blended_weights(mixing_weights, m)
+        for w in ws:
+            depth = self.depth if self.depth > 0 else np.random.randint(1, 4)
+            ops = np.random.choice(self.ops, depth, replace=True)
+            img_aug = img_orig  # no ops are in-place, deep copy not necessary
+            for op in ops:
+                img_aug = op(img_aug)
+            img = Image.blend(img, img_aug, w)
+        return img
+
+    def _apply_basic(self, img, mixing_weights, m):
+        # This is a literal adaptation of the paper/official implementation without normalizations and
+        # PIL <-> Numpy conversions between every op. It is still quite CPU compute heavy compared to the
+        # typical augmentation transforms, could use a GPU / Kornia implementation.
+        img_shape = img.size[0], img.size[1], len(img.getbands())
+        mixed = np.zeros(img_shape, dtype=np.float32)
+        for mw in mixing_weights:
+            depth = self.depth if self.depth > 0 else np.random.randint(1, 4)
+            ops = np.random.choice(self.ops, depth, replace=True)
+            img_aug = img  # no ops are in-place, deep copy not necessary
+            for op in ops:
+                img_aug = op(img_aug)
+            mixed += mw * np.asarray(img_aug, dtype=np.float32)
+        np.clip(mixed, 0, 255., out=mixed)
+        mixed = Image.fromarray(mixed.astype(np.uint8))
+        return Image.blend(img, mixed, m)
+
+    def __call__(self, img):
+        mixing_weights = np.float32(np.random.dirichlet([self.alpha] * self.width))
+        m = np.float32(np.random.beta(self.alpha, self.alpha))
+        if self.blended:
+            mixed = self._apply_blended(img, mixing_weights, m)
+        else:
+            mixed = self._apply_basic(img, mixing_weights, m)
+        return mixed
+
+    def __repr__(self):
+        fs = self.__class__.__name__ + f'(alpha={self.alpha}, width={self.width}, depth={self.depth}, ops='
+        for op in self.ops:
+            fs += f'\n\t{op}'
+        fs += ')'
+        return fs
+
+
+def augment_and_mix_transform(config_str: str, hparams: Optional[Dict] = None):
+    """ Create AugMix PyTorch transform
+
+    Args:
+        config_str (str): String defining configuration of random augmentation. Consists of multiple sections separated
+            by dashes ('-'). The first section defines the specific variant of rand augment (currently only 'rand').
+            The remaining sections, not order sepecific determine
+                'm' - integer magnitude (severity) of augmentation mix (default: 3)
+                'w' - integer width of augmentation chain (default: 3)
+                'd' - integer depth of augmentation chain (-1 is random [1, 3], default: -1)
+                'b' - integer (bool), blend each branch of chain into end result without a final blend, less CPU (default: 0)
+                'mstd' -  float std deviation of magnitude noise applied (default: 0)
+            Ex 'augmix-m5-w4-d2' results in AugMix with severity 5, chain width 4, chain depth 2
+
+        hparams: Other hparams (kwargs) for the Augmentation transforms
+
+    Returns:
+         A PyTorch compatible Transform
+    """
+    magnitude = 3
+    width = 3
+    depth = -1
+    alpha = 1.
+    blended = False
+    config = config_str.split('-')
+    assert config[0] == 'augmix'
+    config = config[1:]
+    for c in config:
+        cs = re.split(r'(\d.*)', c)
+        if len(cs) < 2:
+            continue
+        key, val = cs[:2]
+        if key == 'mstd':
+            # noise param injected via hparams for now
+            hparams.setdefault('magnitude_std', float(val))
+        elif key == 'm':
+            magnitude = int(val)
+        elif key == 'w':
+            width = int(val)
+        elif key == 'd':
+            depth = int(val)
+        elif key == 'a':
+            alpha = float(val)
+        elif key == 'b':
+            blended = bool(val)
+        else:
+            assert False, 'Unknown AugMix config section'
+    hparams.setdefault('magnitude_std', float('inf'))  # default to uniform sampling (if not set via mstd arg)
+    ops = augmix_ops(magnitude=magnitude, hparams=hparams)
+    return AugMixAugment(ops, alpha=alpha, width=width, depth=depth, blended=blended)

timm/data/config.py

@@ -0,0 +1,129 @@
+import logging
+from .constants import *
+
+
+_logger = logging.getLogger(__name__)
+
+
+def resolve_data_config(
+        args=None,
+        pretrained_cfg=None,
+        model=None,
+        use_test_size=False,
+        verbose=False
+):
+    assert model or args or pretrained_cfg, "At least one of model, args, or pretrained_cfg required for data config."
+    args = args or {}
+    pretrained_cfg = pretrained_cfg or {}
+    if not pretrained_cfg and model is not None and hasattr(model, 'pretrained_cfg'):
+        pretrained_cfg = model.pretrained_cfg
+    data_config = {}
+
+    # Resolve input/image size
+    in_chans = 3
+    if args.get('in_chans', None) is not None:
+        in_chans = args['in_chans']
+    elif args.get('chans', None) is not None:
+        in_chans = args['chans']
+
+    input_size = (in_chans, 224, 224)
+    if args.get('input_size', None) is not None:
+        assert isinstance(args['input_size'], (tuple, list))
+        assert len(args['input_size']) == 3
+        input_size = tuple(args['input_size'])
+        in_chans = input_size[0]  # input_size overrides in_chans
+    elif args.get('img_size', None) is not None:
+        assert isinstance(args['img_size'], int)
+        input_size = (in_chans, args['img_size'], args['img_size'])
+    else:
+        if use_test_size and pretrained_cfg.get('test_input_size', None) is not None:
+            input_size = pretrained_cfg['test_input_size']
+        elif pretrained_cfg.get('input_size', None) is not None:
+            input_size = pretrained_cfg['input_size']
+    data_config['input_size'] = input_size
+
+    # resolve interpolation method
+    data_config['interpolation'] = 'bicubic'
+    if args.get('interpolation', None):
+        data_config['interpolation'] = args['interpolation']
+    elif pretrained_cfg.get('interpolation', None):
+        data_config['interpolation'] = pretrained_cfg['interpolation']
+
+    # resolve dataset + model mean for normalization
+    data_config['mean'] = IMAGENET_DEFAULT_MEAN
+    if args.get('mean', None) is not None:
+        mean = tuple(args['mean'])
+        if len(mean) == 1:
+            mean = tuple(list(mean) * in_chans)
+        else:
+            assert len(mean) == in_chans
+        data_config['mean'] = mean
+    elif pretrained_cfg.get('mean', None):
+        data_config['mean'] = pretrained_cfg['mean']
+
+    # resolve dataset + model std deviation for normalization
+    data_config['std'] = IMAGENET_DEFAULT_STD
+    if args.get('std', None) is not None:
+        std = tuple(args['std'])
+        if len(std) == 1:
+            std = tuple(list(std) * in_chans)
+        else:
+            assert len(std) == in_chans
+        data_config['std'] = std
+    elif pretrained_cfg.get('std', None):
+        data_config['std'] = pretrained_cfg['std']
+
+    # resolve default inference crop
+    crop_pct = DEFAULT_CROP_PCT
+    if args.get('crop_pct', None):
+        crop_pct = args['crop_pct']
+    else:
+        if use_test_size and pretrained_cfg.get('test_crop_pct', None):
+            crop_pct = pretrained_cfg['test_crop_pct']
+        elif pretrained_cfg.get('crop_pct', None):
+            crop_pct = pretrained_cfg['crop_pct']
+    data_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 pretrained_cfg.get('crop_mode', None):
+        crop_mode = pretrained_cfg['crop_mode']
+    data_config['crop_mode'] = crop_mode
+
+    if verbose:
+        _logger.info('Data processing configuration for current model + dataset:')
+        for n, v in data_config.items():
+            _logger.info('\t%s: %s' % (n, str(v)))
+
+    return data_config
+
+
+def resolve_model_data_config(
+        model,
+        args=None,
+        pretrained_cfg=None,
+        use_test_size=False,
+        verbose=False,
+):
+    """ Resolve Model Data Config
+    This is equivalent to resolve_data_config() but with arguments re-ordered to put model first.
+
+    Args:
+        model (nn.Module): the model instance
+        args (dict): command line arguments / configuration in dict form (overrides pretrained_cfg)
+        pretrained_cfg (dict): pretrained model config (overrides pretrained_cfg attached to model)
+        use_test_size (bool): use the test time input resolution (if one exists) instead of default train resolution
+        verbose (bool): enable extra logging of resolved values
+
+    Returns:
+        dictionary of config
+    """
+    return resolve_data_config(
+        args=args,
+        pretrained_cfg=pretrained_cfg,
+        model=model,
+        use_test_size=use_test_size,
+        verbose=verbose,
+    )

timm/data/constants.py

@@ -0,0 +1,10 @@
+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)
+IMAGENET_INCEPTION_STD = (0.5, 0.5, 0.5)
+IMAGENET_DPN_MEAN = (124 / 255, 117 / 255, 104 / 255)
+IMAGENET_DPN_STD = tuple([1 / (.0167 * 255)] * 3)
+OPENAI_CLIP_MEAN = (0.48145466, 0.4578275, 0.40821073)
+OPENAI_CLIP_STD = (0.26862954, 0.26130258, 0.27577711)

timm/data/dataset.py

@@ -0,0 +1,200 @@
+""" Quick n Simple Image Folder, Tarfile based DataSet
+
+Hacked together by / Copyright 2019, Ross Wightman
+"""
+import io
+import logging
+from typing import Optional
+
+import torch
+import torch.utils.data as data
+from PIL import Image
+
+from .readers import create_reader
+
+_logger = logging.getLogger(__name__)
+
+
+_ERROR_RETRY = 50
+
+
+class ImageDataset(data.Dataset):
+
+    def __init__(
+            self,
+            root,
+            reader=None,
+            split='train',
+            class_map=None,
+            load_bytes=False,
+            input_img_mode='RGB',
+            transform=None,
+            target_transform=None,
+    ):
+        if reader is None or isinstance(reader, str):
+            reader = create_reader(
+                reader or '',
+                root=root,
+                split=split,
+                class_map=class_map
+            )
+        self.reader = reader
+        self.load_bytes = load_bytes
+        self.input_img_mode = input_img_mode
+        self.transform = transform
+        self.target_transform = target_transform
+        self._consecutive_errors = 0
+
+    def __getitem__(self, index):
+        img, target = self.reader[index]
+
+        try:
+            img = img.read() if self.load_bytes else Image.open(img)
+        except Exception as e:
+            _logger.warning(f'Skipped sample (index {index}, file {self.reader.filename(index)}). {str(e)}')
+            self._consecutive_errors += 1
+            if self._consecutive_errors < _ERROR_RETRY:
+                return self.__getitem__((index + 1) % len(self.reader))
+            else:
+                raise e
+        self._consecutive_errors = 0
+
+        if self.input_img_mode and not self.load_bytes:
+            img = img.convert(self.input_img_mode)
+        if self.transform is not None:
+            img = self.transform(img)
+
+        if target is None:
+            target = -1
+        elif self.target_transform is not None:
+            target = self.target_transform(target)
+
+        return img, target
+
+    def __len__(self):
+        return len(self.reader)
+
+    def filename(self, index, basename=False, absolute=False):
+        return self.reader.filename(index, basename, absolute)
+
+    def filenames(self, basename=False, absolute=False):
+        return self.reader.filenames(basename, absolute)
+
+
+class IterableImageDataset(data.IterableDataset):
+
+    def __init__(
+            self,
+            root,
+            reader=None,
+            split='train',
+            class_map=None,
+            is_training=False,
+            batch_size=1,
+            num_samples=None,
+            seed=42,
+            repeats=0,
+            download=False,
+            input_img_mode='RGB',
+            input_key=None,
+            target_key=None,
+            transform=None,
+            target_transform=None,
+            max_steps=None,
+    ):
+        assert reader is not None
+        if isinstance(reader, str):
+            self.reader = create_reader(
+                reader,
+                root=root,
+                split=split,
+                class_map=class_map,
+                is_training=is_training,
+                batch_size=batch_size,
+                num_samples=num_samples,
+                seed=seed,
+                repeats=repeats,
+                download=download,
+                input_img_mode=input_img_mode,
+                input_key=input_key,
+                target_key=target_key,
+                max_steps=max_steps,
+            )
+        else:
+            self.reader = reader
+        self.transform = transform
+        self.target_transform = target_transform
+        self._consecutive_errors = 0
+
+    def __iter__(self):
+        for img, target in self.reader:
+            if self.transform is not None:
+                img = self.transform(img)
+            if self.target_transform is not None:
+                target = self.target_transform(target)
+            yield img, target
+
+    def __len__(self):
+        if hasattr(self.reader, '__len__'):
+            return len(self.reader)
+        else:
+            return 0
+
+    def set_epoch(self, count):
+        # TFDS and WDS need external epoch count for deterministic cross process shuffle
+        if hasattr(self.reader, 'set_epoch'):
+            self.reader.set_epoch(count)
+
+    def set_loader_cfg(
+            self,
+            num_workers: Optional[int] = None,
+    ):
+        # TFDS and WDS readers need # workers for correct # samples estimate before loader processes created
+        if hasattr(self.reader, 'set_loader_cfg'):
+            self.reader.set_loader_cfg(num_workers=num_workers)
+
+    def filename(self, index, basename=False, absolute=False):
+        assert False, 'Filename lookup by index not supported, use filenames().'
+
+    def filenames(self, basename=False, absolute=False):
+        return self.reader.filenames(basename, absolute)
+
+
+class AugMixDataset(torch.utils.data.Dataset):
+    """Dataset wrapper to perform AugMix or other clean/augmentation mixes"""
+
+    def __init__(self, dataset, num_splits=2):
+        self.augmentation = None
+        self.normalize = None
+        self.dataset = dataset
+        if self.dataset.transform is not None:
+            self._set_transforms(self.dataset.transform)
+        self.num_splits = num_splits
+
+    def _set_transforms(self, x):
+        assert isinstance(x, (list, tuple)) and len(x) == 3, 'Expecting a tuple/list of 3 transforms'
+        self.dataset.transform = x[0]
+        self.augmentation = x[1]
+        self.normalize = x[2]
+
+    @property
+    def transform(self):
+        return self.dataset.transform
+
+    @transform.setter
+    def transform(self, x):
+        self._set_transforms(x)
+
+    def _normalize(self, x):
+        return x if self.normalize is None else self.normalize(x)
+
+    def __getitem__(self, i):
+        x, y = self.dataset[i]  # all splits share the same dataset base transform
+        x_list = [self._normalize(x)]  # first split only normalizes (this is the 'clean' split)
+        # run the full augmentation on the remaining splits
+        for _ in range(self.num_splits - 1):
+            x_list.append(self._normalize(self.augmentation(x)))
+        return tuple(x_list), y
+
+    def __len__(self):
+        return len(self.dataset)

timm/data/dataset_factory.py

@@ -0,0 +1,224 @@
+""" Dataset Factory
+
+Hacked together by / Copyright 2021, Ross Wightman
+"""
+import os
+from typing import Optional
+
+from torchvision.datasets import CIFAR100, CIFAR10, MNIST, KMNIST, FashionMNIST, ImageFolder
+try:
+    from torchvision.datasets import Places365
+    has_places365 = True
+except ImportError:
+    has_places365 = False
+try:
+    from torchvision.datasets import INaturalist
+    has_inaturalist = True
+except ImportError:
+    has_inaturalist = False
+try:
+    from torchvision.datasets import QMNIST
+    has_qmnist = True
+except ImportError:
+    has_qmnist = False
+try:
+    from torchvision.datasets import ImageNet
+    has_imagenet = True
+except ImportError:
+    has_imagenet = False
+
+from .dataset import IterableImageDataset, ImageDataset
+
+_TORCH_BASIC_DS = dict(
+    cifar10=CIFAR10,
+    cifar100=CIFAR100,
+    mnist=MNIST,
+    kmnist=KMNIST,
+    fashion_mnist=FashionMNIST,
+)
+_TRAIN_SYNONYM = dict(train=None, training=None)
+_EVAL_SYNONYM = dict(val=None, valid=None, validation=None, eval=None, evaluation=None)
+
+
+def _search_split(root, split):
+    # look for sub-folder with name of split in root and use that if it exists
+    split_name = split.split('[')[0]
+    try_root = os.path.join(root, split_name)
+    if os.path.exists(try_root):
+        return try_root
+
+    def _try(syn):
+        for s in syn:
+            try_root = os.path.join(root, s)
+            if os.path.exists(try_root):
+                return try_root
+        return root
+    if split_name in _TRAIN_SYNONYM:
+        root = _try(_TRAIN_SYNONYM)
+    elif split_name in _EVAL_SYNONYM:
+        root = _try(_EVAL_SYNONYM)
+    return root
+
+
+def create_dataset(
+        name: str,
+        root: Optional[str] = None,
+        split: str = 'validation',
+        search_split: bool = True,
+        class_map: dict = None,
+        load_bytes: bool = False,
+        is_training: bool = False,
+        download: bool = False,
+        batch_size: int = 1,
+        num_samples: Optional[int] = None,
+        seed: int = 42,
+        repeats: int = 0,
+        input_img_mode: str = 'RGB',
+        **kwargs,
+):
+    """ Dataset factory method
+
+    In parentheses after each arg are the type of dataset supported for each arg, one of:
+      * folder - default, timm folder (or tar) based ImageDataset
+      * torch - torchvision based datasets
+      * HFDS - Hugging Face Datasets
+      * TFDS - Tensorflow-datasets wrapper in IterabeDataset interface via IterableImageDataset
+      * WDS - Webdataset
+      * all - any of the above
+
+    Args:
+        name: dataset name, empty is okay for folder based datasets
+        root: root folder of dataset (all)
+        split: dataset split (all)
+        search_split: search for split specific child fold from root so one can specify
+            `imagenet/` instead of `/imagenet/val`, etc on cmd line / config. (folder, torch/folder)
+        class_map: specify class -> index mapping via text file or dict (folder)
+        load_bytes: load data, return images as undecoded bytes (folder)
+        download: download dataset if not present and supported (HFDS, TFDS, torch)
+        is_training: create dataset in train mode, this is different from the split.
+            For Iterable / TDFS it enables shuffle, ignored for other datasets. (TFDS, WDS)
+        batch_size: batch size hint for (TFDS, WDS)
+        seed: seed for iterable datasets (TFDS, WDS)
+        repeats: dataset repeats per iteration i.e. epoch (TFDS, WDS)
+        input_img_mode: Input image color conversion mode e.g. 'RGB', 'L' (folder, TFDS, WDS, HFDS)
+        **kwargs: other args to pass to dataset
+
+    Returns:
+        Dataset object
+    """
+    kwargs = {k: v for k, v in kwargs.items() if v is not None}
+    name = name.lower()
+    if name.startswith('torch/'):
+        name = name.split('/', 2)[-1]
+        torch_kwargs = dict(root=root, download=download, **kwargs)
+        if name in _TORCH_BASIC_DS:
+            ds_class = _TORCH_BASIC_DS[name]
+            use_train = split in _TRAIN_SYNONYM
+            ds = ds_class(train=use_train, **torch_kwargs)
+        elif name == 'inaturalist' or name == 'inat':
+            assert has_inaturalist, 'Please update to PyTorch 1.10, torchvision 0.11+ for Inaturalist'
+            target_type = 'full'
+            split_split = split.split('/')
+            if len(split_split) > 1:
+                target_type = split_split[0].split('_')
+                if len(target_type) == 1:
+                    target_type = target_type[0]
+                split = split_split[-1]
+            if split in _TRAIN_SYNONYM:
+                split = '2021_train'
+            elif split in _EVAL_SYNONYM:
+                split = '2021_valid'
+            ds = INaturalist(version=split, target_type=target_type, **torch_kwargs)
+        elif name == 'places365':
+            assert has_places365, 'Please update to a newer PyTorch and torchvision for Places365 dataset.'
+            if split in _TRAIN_SYNONYM:
+                split = 'train-standard'
+            elif split in _EVAL_SYNONYM:
+                split = 'val'
+            ds = Places365(split=split, **torch_kwargs)
+        elif name == 'qmnist':
+            assert has_qmnist, 'Please update to a newer PyTorch and torchvision for QMNIST dataset.'
+            use_train = split in _TRAIN_SYNONYM
+            ds = QMNIST(train=use_train, **torch_kwargs)
+        elif name == 'imagenet':
+            assert has_imagenet, 'Please update to a newer PyTorch and torchvision for ImageNet dataset.'
+            if split in _EVAL_SYNONYM:
+                split = 'val'
+            ds = ImageNet(split=split, **torch_kwargs)
+        elif name == 'image_folder' or name == 'folder':
+            # in case torchvision ImageFolder is preferred over timm ImageDataset for some reason
+            if search_split and os.path.isdir(root):
+                # look for split specific sub-folder in root
+                root = _search_split(root, split)
+            ds = ImageFolder(root, **kwargs)
+        else:
+            assert False, f"Unknown torchvision dataset {name}"
+    elif name.startswith('hfds/'):
+        # NOTE right now, HF datasets default arrow format is a random-access Dataset,
+        # There will be a IterableDataset variant too, TBD
+        ds = ImageDataset(
+            root,
+            reader=name,
+            split=split,
+            class_map=class_map,
+            input_img_mode=input_img_mode,
+            **kwargs,
+        )
+    elif name.startswith('hfids/'):
+        ds = IterableImageDataset(
+            root,
+            reader=name,
+            split=split,
+            class_map=class_map,
+            is_training=is_training,
+            download=download,
+            batch_size=batch_size,
+            num_samples=num_samples,
+            repeats=repeats,
+            seed=seed,
+            input_img_mode=input_img_mode,
+            **kwargs
+        )
+    elif name.startswith('tfds/'):
+        ds = IterableImageDataset(
+            root,
+            reader=name,
+            split=split,
+            class_map=class_map,
+            is_training=is_training,
+            download=download,
+            batch_size=batch_size,
+            num_samples=num_samples,
+            repeats=repeats,
+            seed=seed,
+            input_img_mode=input_img_mode,
+            **kwargs
+        )
+    elif name.startswith('wds/'):
+        ds = IterableImageDataset(
+            root,
+            reader=name,
+            split=split,
+            class_map=class_map,
+            is_training=is_training,
+            batch_size=batch_size,
+            num_samples=num_samples,
+            repeats=repeats,
+            seed=seed,
+            input_img_mode=input_img_mode,
+            **kwargs
+        )
+    else:
+        # FIXME support more advance split cfg for ImageFolder/Tar datasets in the future
+        if search_split and os.path.isdir(root):
+            # look for split specific sub-folder in root
+            root = _search_split(root, split)
+        ds = ImageDataset(
+            root,
+            reader=name,
+            class_map=class_map,
+            load_bytes=load_bytes,
+            input_img_mode=input_img_mode,
+            **kwargs,
+        )
+    return ds

timm/data/dataset_info.py

@@ -0,0 +1,73 @@
+from abc import ABC, abstractmethod
+from typing import Dict, List, Optional, Union
+
+
+class DatasetInfo(ABC):
+
+    def __init__(self):
+        pass
+
+    @abstractmethod
+    def num_classes(self):
+        pass
+
+    @abstractmethod
+    def label_names(self):
+        pass
+
+    @abstractmethod
+    def label_descriptions(self, detailed: bool = False, as_dict: bool = False) -> Union[List[str], Dict[str, str]]:
+        pass
+
+    @abstractmethod
+    def index_to_label_name(self, index) -> str:
+        pass
+
+    @abstractmethod
+    def index_to_description(self, index: int, detailed: bool = False) -> str:
+        pass
+
+    @abstractmethod
+    def label_name_to_description(self, label: str, detailed: bool = False) -> str:
+        pass
+
+
+class CustomDatasetInfo(DatasetInfo):
+    """ DatasetInfo that wraps passed values for custom datasets."""
+
+    def __init__(
+            self,
+            label_names: Union[List[str], Dict[int, str]],
+            label_descriptions: Optional[Dict[str, str]] = None
+    ):
+        super().__init__()
+        assert len(label_names) > 0
+        self._label_names = label_names  # label index => label name mapping
+        self._label_descriptions = label_descriptions  # label name => label description mapping
+        if self._label_descriptions is not None:
+            # validate descriptions (label names required)
+            assert isinstance(self._label_descriptions, dict)
+            for n in self._label_names:
+                assert n in self._label_descriptions
+
+    def num_classes(self):
+        return len(self._label_names)
+
+    def label_names(self):
+        return self._label_names
+
+    def label_descriptions(self, detailed: bool = False, as_dict: bool = False) -> Union[List[str], Dict[str, str]]:
+        return self._label_descriptions
+
+    def label_name_to_description(self, label: str, detailed: bool = False) -> str:
+        if self._label_descriptions:
+            return self._label_descriptions[label]
+        return label  # return label name itself if a descriptions is not present
+
+    def index_to_label_name(self, index) -> str:
+        assert 0 <= index < len(self._label_names)
+        return self._label_names[index]
+
+    def index_to_description(self, index: int, detailed: bool = False) -> str:
+        label = self.index_to_label_name(index)
+        return self.label_name_to_description(label, detailed=detailed)

timm/data/distributed_sampler.py

@@ -0,0 +1,135 @@
+import math
+import torch
+from torch.utils.data import Sampler
+import torch.distributed as dist
+
+
+class OrderedDistributedSampler(Sampler):
+    """Sampler that restricts data loading to a subset of the dataset.
+    It is especially useful in conjunction with
+    :class:`torch.nn.parallel.DistributedDataParallel`. In such case, each
+    process can pass a DistributedSampler instance as a DataLoader sampler,
+    and load a subset of the original dataset that is exclusive to it.
+    .. note::
+        Dataset is assumed to be of constant size.
+    Arguments:
+        dataset: Dataset used for sampling.
+        num_replicas (optional): Number of processes participating in
+            distributed training.
+        rank (optional): Rank of the current process within num_replicas.
+    """
+
+    def __init__(self, dataset, num_replicas=None, rank=None):
+        if num_replicas is None:
+            if not dist.is_available():
+                raise RuntimeError("Requires distributed package to be available")
+            num_replicas = dist.get_world_size()
+        if rank is None:
+            if not dist.is_available():
+                raise RuntimeError("Requires distributed package to be available")
+            rank = dist.get_rank()
+        self.dataset = dataset
+        self.num_replicas = num_replicas
+        self.rank = rank
+        self.num_samples = int(math.ceil(len(self.dataset) * 1.0 / self.num_replicas))
+        self.total_size = self.num_samples * self.num_replicas
+
+    def __iter__(self):
+        indices = list(range(len(self.dataset)))
+
+        # add extra samples to make it evenly divisible
+        indices += indices[:(self.total_size - len(indices))]
+        assert len(indices) == self.total_size
+
+        # subsample
+        indices = indices[self.rank:self.total_size:self.num_replicas]
+        assert len(indices) == self.num_samples
+
+        return iter(indices)
+
+    def __len__(self):
+        return self.num_samples
+
+
+class RepeatAugSampler(Sampler):
+    """Sampler that restricts data loading to a subset of the dataset for distributed,
+    with repeated augmentation.
+    It ensures that different each augmented version of a sample will be visible to a
+    different process (GPU). Heavily based on torch.utils.data.DistributedSampler
+
+    This sampler was taken from https://github.com/facebookresearch/deit/blob/0c4b8f60/samplers.py
+    Used in
+    Copyright (c) 2015-present, Facebook, Inc.
+    """
+
+    def __init__(
+            self,
+            dataset,
+            num_replicas=None,
+            rank=None,
+            shuffle=True,
+            num_repeats=3,
+            selected_round=256,
+            selected_ratio=0,
+    ):
+        if num_replicas is None:
+            if not dist.is_available():
+                raise RuntimeError("Requires distributed package to be available")
+            num_replicas = dist.get_world_size()
+        if rank is None:
+            if not dist.is_available():
+                raise RuntimeError("Requires distributed package to be available")
+            rank = dist.get_rank()
+        self.dataset = dataset
+        self.num_replicas = num_replicas
+        self.rank = rank
+        self.shuffle = shuffle
+        self.num_repeats = num_repeats
+        self.epoch = 0
+        self.num_samples = int(math.ceil(len(self.dataset) * num_repeats / self.num_replicas))
+        self.total_size = self.num_samples * self.num_replicas
+        # Determine the number of samples to select per epoch for each rank.
+        # num_selected logic defaults to be the same as original RASampler impl, but this one can be tweaked
+        # via selected_ratio and selected_round args.
+        selected_ratio = selected_ratio or num_replicas  # ratio to reduce selected samples by, num_replicas if 0
+        if selected_round:
+            self.num_selected_samples = int(math.floor(
+                 len(self.dataset) // selected_round * selected_round / selected_ratio))
+        else:
+            self.num_selected_samples = int(math.ceil(len(self.dataset) / selected_ratio))
+
+    def __iter__(self):
+        # deterministically shuffle based on epoch
+        g = torch.Generator()
+        g.manual_seed(self.epoch)
+        if self.shuffle:
+            indices = torch.randperm(len(self.dataset), generator=g)
+        else:
+            indices = torch.arange(start=0, end=len(self.dataset))
+
+        # produce repeats e.g. [0, 0, 0, 1, 1, 1, 2, 2, 2....]
+        if isinstance(self.num_repeats, float) and not self.num_repeats.is_integer():
+            # resample for repeats w/ non-integer ratio
+            repeat_size = math.ceil(self.num_repeats * len(self.dataset))
+            indices = indices[torch.tensor([int(i // self.num_repeats) for i in range(repeat_size)])]
+        else:
+            indices = torch.repeat_interleave(indices, repeats=int(self.num_repeats), dim=0)
+        indices = indices.tolist()  # leaving as tensor thrashes dataloader memory
+        # add extra samples to make it evenly divisible
+        padding_size = self.total_size - len(indices)
+        if padding_size > 0:
+            indices += indices[:padding_size]
+        assert len(indices) == self.total_size
+
+        # subsample per rank
+        indices = indices[self.rank:self.total_size:self.num_replicas]
+        assert len(indices) == self.num_samples
+
+        # return up to num selected samples
+        return iter(indices[:self.num_selected_samples])
+
+    def __len__(self):
+        return self.num_selected_samples
+
+    def set_epoch(self, epoch):
+        self.epoch = epoch

timm/data/loader.py

@@ -0,0 +1,413 @@
+""" Loader Factory, Fast Collate, CUDA Prefetcher
+
+Prefetcher and Fast Collate inspired by NVIDIA APEX example at
+https://github.com/NVIDIA/apex/commit/d5e2bb4bdeedd27b1dfaf5bb2b24d6c000dee9be#diff-cf86c282ff7fba81fad27a559379d5bf
+
+Hacked together by / Copyright 2019, Ross Wightman
+"""
+
+# Copyright 2026 Kiel University
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+# Based on pytorch-image-models (timm); see NOTICE.
+#
+# Modifications:
+# - Increased the DataLoader prefetch factor for ProgResViT training.
+
+import logging
+import random
+from contextlib import suppress
+from functools import partial
+from itertools import repeat
+from typing import Callable, Optional, Tuple, Union
+
+import torch
+import torch.utils.data
+import numpy as np
+
+from .constants import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD
+from .dataset import IterableImageDataset, ImageDataset
+from .distributed_sampler import OrderedDistributedSampler, RepeatAugSampler
+from .random_erasing import RandomErasing
+from .mixup import FastCollateMixup
+from .transforms_factory import create_transform
+
+_logger = logging.getLogger(__name__)
+
+
+def fast_collate(batch):
+    """ A fast collation function optimized for uint8 images (np array or torch) and int64 targets (labels)"""
+    assert isinstance(batch[0], tuple)
+    batch_size = len(batch)
+    if isinstance(batch[0][0], tuple):
+        # This branch 'deinterleaves' and flattens tuples of input tensors into one tensor ordered by position
+        # such that all tuple of position n will end up in a torch.split(tensor, batch_size) in nth position
+        inner_tuple_size = len(batch[0][0])
+        flattened_batch_size = batch_size * inner_tuple_size
+        targets = torch.zeros(flattened_batch_size, dtype=torch.int64)
+        tensor = torch.zeros((flattened_batch_size, *batch[0][0][0].shape), dtype=torch.uint8)
+        for i in range(batch_size):
+            assert len(batch[i][0]) == inner_tuple_size  # all input tensor tuples must be same length
+            for j in range(inner_tuple_size):
+                targets[i + j * batch_size] = batch[i][1]
+                tensor[i + j * batch_size] += torch.from_numpy(batch[i][0][j])
+        return tensor, targets
+    elif isinstance(batch[0][0], np.ndarray):
+        targets = torch.tensor([b[1] for b in batch], dtype=torch.int64)
+        assert len(targets) == batch_size
+        tensor = torch.zeros((batch_size, *batch[0][0].shape), dtype=torch.uint8)
+        for i in range(batch_size):
+            tensor[i] += torch.from_numpy(batch[i][0])
+        return tensor, targets
+    elif isinstance(batch[0][0], torch.Tensor):
+        targets = torch.tensor([b[1] for b in batch], dtype=torch.int64)
+        assert len(targets) == batch_size
+        tensor = torch.zeros((batch_size, *batch[0][0].shape), dtype=torch.uint8)
+        for i in range(batch_size):
+            tensor[i].copy_(batch[i][0])
+        return tensor, targets
+    else:
+        assert False
+
+
+def adapt_to_chs(x, n):
+    if not isinstance(x, (tuple, list)):
+        x = tuple(repeat(x, n))
+    elif len(x) != n:
+        x_mean = np.mean(x).item()
+        x = (x_mean,) * n
+        _logger.warning(f'Pretrained mean/std different shape than model, using avg value {x}.')
+    else:
+        assert len(x) == n, 'normalization stats must match image channels'
+    return x
+
+
+class PrefetchLoader:
+
+    def __init__(
+            self,
+            loader,
+            mean=IMAGENET_DEFAULT_MEAN,
+            std=IMAGENET_DEFAULT_STD,
+            channels=3,
+            device=torch.device('cuda'),
+            img_dtype=torch.float32,
+            fp16=False,
+            re_prob=0.,
+            re_mode='const',
+            re_count=1,
+            re_num_splits=0):
+
+        mean = adapt_to_chs(mean, channels)
+        std = adapt_to_chs(std, channels)
+        normalization_shape = (1, channels, 1, 1)
+
+        self.loader = loader
+        self.device = device
+        if fp16:
+            # fp16 arg is deprecated, but will override dtype arg if set for bwd compat
+            img_dtype = torch.float16
+        self.img_dtype = img_dtype
+        self.mean = torch.tensor(
+            [x * 255 for x in mean], device=device, dtype=img_dtype).view(normalization_shape)
+        self.std = torch.tensor(
+            [x * 255 for x in std], device=device, dtype=img_dtype).view(normalization_shape)
+        if re_prob > 0.:
+            self.random_erasing = RandomErasing(
+                probability=re_prob,
+                mode=re_mode,
+                max_count=re_count,
+                num_splits=re_num_splits,
+                device=device,
+            )
+        else:
+            self.random_erasing = None
+        self.is_cuda = torch.cuda.is_available() and device.type == 'cuda'
+
+    def __iter__(self):
+        first = True
+        if self.is_cuda:
+            stream = torch.cuda.Stream()
+            stream_context = partial(torch.cuda.stream, stream=stream)
+        else:
+            stream = None
+            stream_context = suppress
+
+        for next_input, next_target in self.loader:
+
+            with stream_context():
+                next_input = next_input.to(device=self.device, non_blocking=True)
+                next_target = next_target.to(device=self.device, non_blocking=True)
+                next_input = next_input.to(self.img_dtype).sub_(self.mean).div_(self.std)
+                if self.random_erasing is not None:
+                    next_input = self.random_erasing(next_input)
+
+            if not first:
+                yield input, target
+            else:
+                first = False
+
+            if stream is not None:
+                torch.cuda.current_stream().wait_stream(stream)
+
+            input = next_input
+            target = next_target
+
+        yield input, target
+
+    def __len__(self):
+        return len(self.loader)
+
+    @property
+    def sampler(self):
+        return self.loader.sampler
+
+    @property
+    def dataset(self):
+        return self.loader.dataset
+
+    @property
+    def mixup_enabled(self):
+        if isinstance(self.loader.collate_fn, FastCollateMixup):
+            return self.loader.collate_fn.mixup_enabled
+        else:
+            return False
+
+    @mixup_enabled.setter
+    def mixup_enabled(self, x):
+        if isinstance(self.loader.collate_fn, FastCollateMixup):
+            self.loader.collate_fn.mixup_enabled = x
+
+
+def _worker_init(worker_id, worker_seeding='all'):
+    worker_info = torch.utils.data.get_worker_info()
+    assert worker_info.id == worker_id
+    if isinstance(worker_seeding, Callable):
+        seed = worker_seeding(worker_info)
+        random.seed(seed)
+        torch.manual_seed(seed)
+        np.random.seed(seed % (2 ** 32 - 1))
+    else:
+        assert worker_seeding in ('all', 'part')
+        # random / torch seed already called in dataloader iter class w/ worker_info.seed
+        # to reproduce some old results (same seed + hparam combo), partial seeding is required (skip numpy re-seed)
+        if worker_seeding == 'all':
+            np.random.seed(worker_info.seed % (2 ** 32 - 1))
+
+
+def create_loader(
+        dataset: Union[ImageDataset, IterableImageDataset],
+        input_size: Union[int, Tuple[int, int], Tuple[int, int, int]],
+        batch_size: int,
+        is_training: bool = False,
+        no_aug: bool = False,
+        re_prob: float = 0.,
+        re_mode: str = 'const',
+        re_count: int = 1,
+        re_split: bool = False,
+        train_crop_mode: Optional[str] = None,
+        scale: Optional[Tuple[float, float]] = None,
+        ratio: Optional[Tuple[float, float]] = None,
+        hflip: float = 0.5,
+        vflip: float = 0.,
+        color_jitter: float = 0.4,
+        color_jitter_prob: Optional[float] = None,
+        grayscale_prob: float = 0.,
+        gaussian_blur_prob: float = 0.,
+        auto_augment: Optional[str] = None,
+        num_aug_repeats: int = 0,
+        num_aug_splits: int = 0,
+        interpolation: str = 'bilinear',
+        mean: Tuple[float, ...] = IMAGENET_DEFAULT_MEAN,
+        std: Tuple[float, ...] = IMAGENET_DEFAULT_STD,
+        num_workers: int = 1,
+        distributed: bool = False,
+        crop_pct: Optional[float] = None,
+        crop_mode: Optional[str] = None,
+        crop_border_pixels: Optional[int] = None,
+        collate_fn: Optional[Callable] = None,
+        pin_memory: bool = False,
+        fp16: bool = False,  # deprecated, use img_dtype
+        img_dtype: torch.dtype = torch.float32,
+        device: torch.device = torch.device('cuda'),
+        use_prefetcher: bool = True,
+        use_multi_epochs_loader: bool = False,
+        persistent_workers: bool = True,
+        worker_seeding: str = 'all',
+        tf_preprocessing: bool = False,
+):
+    """
+
+    Args:
+        dataset: The image dataset to load.
+        input_size: Target input size (channels, height, width) tuple or size scalar.
+        batch_size: Number of samples in a batch.
+        is_training: Return training (random) transforms.
+        no_aug: Disable augmentation for training (useful for debug).
+        re_prob: Random erasing probability.
+        re_mode: Random erasing fill mode.
+        re_count: Number of random erasing regions.
+        re_split: Control split of random erasing across batch size.
+        scale: Random resize scale range (crop area, < 1.0 => zoom in).
+        ratio: Random aspect ratio range (crop ratio for RRC, ratio adjustment factor for RKR).
+        hflip: Horizontal flip probability.
+        vflip: Vertical flip probability.
+        color_jitter: Random color jitter component factors (brightness, contrast, saturation, hue).
+            Scalar is applied as (scalar,) * 3 (no hue).
+        color_jitter_prob: Apply color jitter with this probability if not None (for SimlCLR-like aug
+        grayscale_prob: Probability of converting image to grayscale (for SimCLR-like aug).
+        gaussian_blur_prob: Probability of applying gaussian blur (for SimCLR-like aug).
+        auto_augment: Auto augment configuration string (see auto_augment.py).
+        num_aug_repeats: Enable special sampler to repeat same augmentation across distributed GPUs.
+        num_aug_splits: Enable mode where augmentations can be split across the batch.
+        interpolation: Image interpolation mode.
+        mean: Image normalization mean.
+        std: Image normalization standard deviation.
+        num_workers: Num worker processes per DataLoader.
+        distributed: Enable dataloading for distributed training.
+        crop_pct: Inference crop percentage (output size / resize size).
+        crop_mode: Inference crop mode. One of ['squash', 'border', 'center']. Defaults to 'center' when None.
+        crop_border_pixels: Inference crop border of specified # pixels around edge of original image.
+        collate_fn: Override default collate_fn.
+        pin_memory: Pin memory for device transfer.
+        fp16: Deprecated argument for half-precision input dtype. Use img_dtype.
+        img_dtype: Data type for input image.
+        device: Device to transfer inputs and targets to.
+        use_prefetcher: Use efficient pre-fetcher to load samples onto device.
+        use_multi_epochs_loader:
+        persistent_workers: Enable persistent worker processes.
+        worker_seeding: Control worker random seeding at init.
+        tf_preprocessing: Use TF 1.0 inference preprocessing for testing model ports.
+
+    Returns:
+        DataLoader
+    """
+    re_num_splits = 0
+    if re_split:
+        # apply RE to second half of batch if no aug split otherwise line up with aug split
+        re_num_splits = num_aug_splits or 2
+    dataset.transform = create_transform(
+        input_size,
+        is_training=is_training,
+        no_aug=no_aug,
+        train_crop_mode=train_crop_mode,
+        scale=scale,
+        ratio=ratio,
+        hflip=hflip,
+        vflip=vflip,
+        color_jitter=color_jitter,
+        color_jitter_prob=color_jitter_prob,
+        grayscale_prob=grayscale_prob,
+        gaussian_blur_prob=gaussian_blur_prob,
+        auto_augment=auto_augment,
+        interpolation=interpolation,
+        mean=mean,
+        std=std,
+        crop_pct=crop_pct,
+        crop_mode=crop_mode,
+        crop_border_pixels=crop_border_pixels,
+        re_prob=re_prob,
+        re_mode=re_mode,
+        re_count=re_count,
+        re_num_splits=re_num_splits,
+        tf_preprocessing=tf_preprocessing,
+        use_prefetcher=use_prefetcher,
+        separate=num_aug_splits > 0,
+    )
+
+    if isinstance(dataset, IterableImageDataset):
+        # give Iterable datasets early knowledge of num_workers so that sample estimates
+        # are correct before worker processes are launched
+        dataset.set_loader_cfg(num_workers=num_workers)
+
+    sampler = None
+    if distributed and not isinstance(dataset, torch.utils.data.IterableDataset):
+        if is_training:
+            if num_aug_repeats:
+                sampler = RepeatAugSampler(dataset, num_repeats=num_aug_repeats)
+            else:
+                sampler = torch.utils.data.distributed.DistributedSampler(dataset)
+        else:
+            # This will add extra duplicate entries to result in equal num
+            # of samples per-process, will slightly alter validation results
+            sampler = OrderedDistributedSampler(dataset)
+    else:
+        assert num_aug_repeats == 0, "RepeatAugment not currently supported in non-distributed or IterableDataset use"
+
+    if collate_fn is None:
+        collate_fn = fast_collate if use_prefetcher else torch.utils.data.dataloader.default_collate
+
+    loader_class = torch.utils.data.DataLoader
+    if use_multi_epochs_loader:
+        loader_class = MultiEpochsDataLoader
+
+    loader_args = dict(
+        batch_size=batch_size,
+        shuffle=not isinstance(dataset, torch.utils.data.IterableDataset) and sampler is None and is_training,
+        num_workers=num_workers,
+        sampler=sampler,
+        collate_fn=collate_fn,
+        prefetch_factor=20,
+        pin_memory=pin_memory,
+        drop_last=is_training,
+        worker_init_fn=partial(_worker_init, worker_seeding=worker_seeding),
+        persistent_workers=persistent_workers
+    )
+    try:
+        loader = loader_class(dataset, **loader_args)
+    except TypeError as e:
+        loader_args.pop('persistent_workers')  # only in Pytorch 1.7+
+        loader = loader_class(dataset, **loader_args)
+    if use_prefetcher:
+        prefetch_re_prob = re_prob if is_training and not no_aug else 0.
+        loader = PrefetchLoader(
+            loader,
+            mean=mean,
+            std=std,
+            channels=input_size[0],
+            device=device,
+            fp16=fp16,  # deprecated, use img_dtype
+            img_dtype=img_dtype,
+            re_prob=prefetch_re_prob,
+            re_mode=re_mode,
+            re_count=re_count,
+            re_num_splits=re_num_splits
+        )
+
+    return loader
+
+
+class MultiEpochsDataLoader(torch.utils.data.DataLoader):
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self._DataLoader__initialized = False
+        if self.batch_sampler is None:
+            self.sampler = _RepeatSampler(self.sampler)
+        else:
+            self.batch_sampler = _RepeatSampler(self.batch_sampler)
+        self._DataLoader__initialized = True
+        self.iterator = super().__iter__()
+
+    def __len__(self):
+        return len(self.sampler) if self.batch_sampler is None else len(self.batch_sampler.sampler)
+
+    def __iter__(self):
+        for i in range(len(self)):
+            yield next(self.iterator)
+
+
+class _RepeatSampler(object):
+    """ Sampler that repeats forever.
+
+    Args:
+        sampler (Sampler)
+    """
+
+    def __init__(self, sampler):
+        self.sampler = sampler
+
+    def __iter__(self):
+        while True:
+            yield from iter(self.sampler)

timm/data/mixup.py

@@ -0,0 +1,315 @@
+""" Mixup and Cutmix
+
+Papers:
+mixup: Beyond Empirical Risk Minimization (https://arxiv.org/abs/1710.09412)
+
+CutMix: Regularization Strategy to Train Strong Classifiers with Localizable Features (https://arxiv.org/abs/1905.04899)
+
+Code Reference:
+CutMix: https://github.com/clovaai/CutMix-PyTorch
+
+Hacked together by / Copyright 2019, Ross Wightman
+"""
+import numpy as np
+import torch
+
+
+def one_hot(x, num_classes, on_value=1., off_value=0.):
+    x = x.long().view(-1, 1)
+    return torch.full((x.size()[0], num_classes), off_value, device=x.device).scatter_(1, x, on_value)
+
+
+def mixup_target(target, num_classes, lam=1., smoothing=0.0):
+    off_value = smoothing / num_classes
+    on_value = 1. - smoothing + off_value
+    y1 = one_hot(target, num_classes, on_value=on_value, off_value=off_value)
+    y2 = one_hot(target.flip(0), num_classes, on_value=on_value, off_value=off_value)
+    return y1 * lam + y2 * (1. - lam)
+
+
+def rand_bbox(img_shape, lam, margin=0., count=None):
+    """ Standard CutMix bounding-box
+    Generates a random square bbox based on lambda value. This impl includes
+    support for enforcing a border margin as percent of bbox dimensions.
+
+    Args:
+        img_shape (tuple): Image shape as tuple
+        lam (float): Cutmix lambda value
+        margin (float): Percentage of bbox dimension to enforce as margin (reduce amount of box outside image)
+        count (int): Number of bbox to generate
+    """
+    ratio = np.sqrt(1 - lam)
+    img_h, img_w = img_shape[-2:]
+    cut_h, cut_w = int(img_h * ratio), int(img_w * ratio)
+    margin_y, margin_x = int(margin * cut_h), int(margin * cut_w)
+    cy = np.random.randint(0 + margin_y, img_h - margin_y, size=count)
+    cx = np.random.randint(0 + margin_x, img_w - margin_x, size=count)
+    yl = np.clip(cy - cut_h // 2, 0, img_h)
+    yh = np.clip(cy + cut_h // 2, 0, img_h)
+    xl = np.clip(cx - cut_w // 2, 0, img_w)
+    xh = np.clip(cx + cut_w // 2, 0, img_w)
+    return yl, yh, xl, xh
+
+
+def rand_bbox_minmax(img_shape, minmax, count=None):
+    """ Min-Max CutMix bounding-box
+    Inspired by Darknet cutmix impl, generates a random rectangular bbox
+    based on min/max percent values applied to each dimension of the input image.
+
+    Typical defaults for minmax are usually in the  .2-.3 for min and .8-.9 range for max.
+
+    Args:
+        img_shape (tuple): Image shape as tuple
+        minmax (tuple or list): Min and max bbox ratios (as percent of image size)
+        count (int): Number of bbox to generate
+    """
+    assert len(minmax) == 2
+    img_h, img_w = img_shape[-2:]
+    cut_h = np.random.randint(int(img_h * minmax[0]), int(img_h * minmax[1]), size=count)
+    cut_w = np.random.randint(int(img_w * minmax[0]), int(img_w * minmax[1]), size=count)
+    yl = np.random.randint(0, img_h - cut_h, size=count)
+    xl = np.random.randint(0, img_w - cut_w, size=count)
+    yu = yl + cut_h
+    xu = xl + cut_w
+    return yl, yu, xl, xu
+
+
+def cutmix_bbox_and_lam(img_shape, lam, ratio_minmax=None, correct_lam=True, count=None):
+    """ Generate bbox and apply lambda correction.
+    """
+    if ratio_minmax is not None:
+        yl, yu, xl, xu = rand_bbox_minmax(img_shape, ratio_minmax, count=count)
+    else:
+        yl, yu, xl, xu = rand_bbox(img_shape, lam, count=count)
+    if correct_lam or ratio_minmax is not None:
+        bbox_area = (yu - yl) * (xu - xl)
+        lam = 1. - bbox_area / float(img_shape[-2] * img_shape[-1])
+    return (yl, yu, xl, xu), lam
+
+
+class Mixup:
+    """ Mixup/Cutmix that applies different params to each element or whole batch
+
+    Args:
+        mixup_alpha (float): mixup alpha value, mixup is active if > 0.
+        cutmix_alpha (float): cutmix alpha value, cutmix is active if > 0.
+        cutmix_minmax (List[float]): cutmix min/max image ratio, cutmix is active and uses this vs alpha if not None.
+        prob (float): probability of applying mixup or cutmix per batch or element
+        switch_prob (float): probability of switching to cutmix instead of mixup when both are active
+        mode (str): how to apply mixup/cutmix params (per 'batch', 'pair' (pair of elements), 'elem' (element)
+        correct_lam (bool): apply lambda correction when cutmix bbox clipped by image borders
+        label_smoothing (float): apply label smoothing to the mixed target tensor
+        num_classes (int): number of classes for target
+    """
+    def __init__(self, mixup_alpha=1., cutmix_alpha=0., cutmix_minmax=None, prob=1.0, switch_prob=0.5,
+                 mode='batch', correct_lam=True, label_smoothing=0.1, num_classes=1000):
+        self.mixup_alpha = mixup_alpha
+        self.cutmix_alpha = cutmix_alpha
+        self.cutmix_minmax = cutmix_minmax
+        if self.cutmix_minmax is not None:
+            assert len(self.cutmix_minmax) == 2
+            # force cutmix alpha == 1.0 when minmax active to keep logic simple & safe
+            self.cutmix_alpha = 1.0
+        self.mix_prob = prob
+        self.switch_prob = switch_prob
+        self.label_smoothing = label_smoothing
+        self.num_classes = num_classes
+        self.mode = mode
+        self.correct_lam = correct_lam  # correct lambda based on clipped area for cutmix
+        self.mixup_enabled = True  # set to false to disable mixing (intended tp be set by train loop)
+
+    def _params_per_elem(self, batch_size):
+        lam = np.ones(batch_size, dtype=np.float32)
+        use_cutmix = np.zeros(batch_size, dtype=bool)
+        if self.mixup_enabled:
+            if self.mixup_alpha > 0. and self.cutmix_alpha > 0.:
+                use_cutmix = np.random.rand(batch_size) < self.switch_prob
+                lam_mix = np.where(
+                    use_cutmix,
+                    np.random.beta(self.cutmix_alpha, self.cutmix_alpha, size=batch_size),
+                    np.random.beta(self.mixup_alpha, self.mixup_alpha, size=batch_size))
+            elif self.mixup_alpha > 0.:
+                lam_mix = np.random.beta(self.mixup_alpha, self.mixup_alpha, size=batch_size)
+            elif self.cutmix_alpha > 0.:
+                use_cutmix = np.ones(batch_size, dtype=bool)
+                lam_mix = np.random.beta(self.cutmix_alpha, self.cutmix_alpha, size=batch_size)
+            else:
+                assert False, "One of mixup_alpha > 0., cutmix_alpha > 0., cutmix_minmax not None should be true."
+            lam = np.where(np.random.rand(batch_size) < self.mix_prob, lam_mix.astype(np.float32), lam)
+        return lam, use_cutmix
+
+    def _params_per_batch(self):
+        lam = 1.
+        use_cutmix = False
+        if self.mixup_enabled and np.random.rand() < self.mix_prob:
+            if self.mixup_alpha > 0. and self.cutmix_alpha > 0.:
+                use_cutmix = np.random.rand() < self.switch_prob
+                lam_mix = np.random.beta(self.cutmix_alpha, self.cutmix_alpha) if use_cutmix else \
+                    np.random.beta(self.mixup_alpha, self.mixup_alpha)
+            elif self.mixup_alpha > 0.:
+                lam_mix = np.random.beta(self.mixup_alpha, self.mixup_alpha)
+            elif self.cutmix_alpha > 0.:
+                use_cutmix = True
+                lam_mix = np.random.beta(self.cutmix_alpha, self.cutmix_alpha)
+            else:
+                assert False, "One of mixup_alpha > 0., cutmix_alpha > 0., cutmix_minmax not None should be true."
+            lam = float(lam_mix)
+        return lam, use_cutmix
+
+    def _mix_elem(self, x):
+        batch_size = len(x)
+        lam_batch, use_cutmix = self._params_per_elem(batch_size)
+        x_orig = x.clone()  # need to keep an unmodified original for mixing source
+        for i in range(batch_size):
+            j = batch_size - i - 1
+            lam = lam_batch[i]
+            if lam != 1.:
+                if use_cutmix[i]:
+                    (yl, yh, xl, xh), lam = cutmix_bbox_and_lam(
+                        x[i].shape, lam, ratio_minmax=self.cutmix_minmax, correct_lam=self.correct_lam)
+                    x[i][:, yl:yh, xl:xh] = x_orig[j][:, yl:yh, xl:xh]
+                    lam_batch[i] = lam
+                else:
+                    x[i] = x[i] * lam + x_orig[j] * (1 - lam)
+        return torch.tensor(lam_batch, device=x.device, dtype=x.dtype).unsqueeze(1)
+
+    def _mix_pair(self, x):
+        batch_size = len(x)
+        lam_batch, use_cutmix = self._params_per_elem(batch_size // 2)
+        x_orig = x.clone()  # need to keep an unmodified original for mixing source
+        for i in range(batch_size // 2):
+            j = batch_size - i - 1
+            lam = lam_batch[i]
+            if lam != 1.:
+                if use_cutmix[i]:
+                    (yl, yh, xl, xh), lam = cutmix_bbox_and_lam(
+                        x[i].shape, lam, ratio_minmax=self.cutmix_minmax, correct_lam=self.correct_lam)
+                    x[i][:, yl:yh, xl:xh] = x_orig[j][:, yl:yh, xl:xh]
+                    x[j][:, yl:yh, xl:xh] = x_orig[i][:, yl:yh, xl:xh]
+                    lam_batch[i] = lam
+                else:
+                    x[i] = x[i] * lam + x_orig[j] * (1 - lam)
+                    x[j] = x[j] * lam + x_orig[i] * (1 - lam)
+        lam_batch = np.concatenate((lam_batch, lam_batch[::-1]))
+        return torch.tensor(lam_batch, device=x.device, dtype=x.dtype).unsqueeze(1)
+
+    def _mix_batch(self, x):
+        lam, use_cutmix = self._params_per_batch()
+        if lam == 1.:
+            return 1.
+        if use_cutmix:
+            (yl, yh, xl, xh), lam = cutmix_bbox_and_lam(
+                x.shape, lam, ratio_minmax=self.cutmix_minmax, correct_lam=self.correct_lam)
+            x[:, :, yl:yh, xl:xh] = x.flip(0)[:, :, yl:yh, xl:xh]
+        else:
+            x_flipped = x.flip(0).mul_(1. - lam)
+            x.mul_(lam).add_(x_flipped)
+        return lam
+
+    def __call__(self, x, target):
+        assert len(x) % 2 == 0, 'Batch size should be even when using this'
+        if self.mode == 'elem':
+            lam = self._mix_elem(x)
+        elif self.mode == 'pair':
+            lam = self._mix_pair(x)
+        else:
+            lam = self._mix_batch(x)
+        target = mixup_target(target, self.num_classes, lam, self.label_smoothing)
+        return x, target
+
+
+class FastCollateMixup(Mixup):
+    """ Fast Collate w/ Mixup/Cutmix that applies different params to each element or whole batch
+
+    A Mixup impl that's performed while collating the batches.
+    """
+
+    def _mix_elem_collate(self, output, batch, half=False):
+        batch_size = len(batch)
+        num_elem = batch_size // 2 if half else batch_size
+        assert len(output) == num_elem
+        lam_batch, use_cutmix = self._params_per_elem(num_elem)
+        for i in range(num_elem):
+            j = batch_size - i - 1
+            lam = lam_batch[i]
+            mixed = batch[i][0]
+            if lam != 1.:
+                if use_cutmix[i]:
+                    if not half:
+                        mixed = mixed.copy()
+                    (yl, yh, xl, xh), lam = cutmix_bbox_and_lam(
+                        output.shape, lam, ratio_minmax=self.cutmix_minmax, correct_lam=self.correct_lam)
+                    mixed[:, yl:yh, xl:xh] = batch[j][0][:, yl:yh, xl:xh]
+                    lam_batch[i] = lam
+                else:
+                    mixed = mixed.astype(np.float32) * lam + batch[j][0].astype(np.float32) * (1 - lam)
+                    np.rint(mixed, out=mixed)
+            output[i] += torch.from_numpy(mixed.astype(np.uint8))
+        if half:
+            lam_batch = np.concatenate((lam_batch, np.ones(num_elem)))
+        return torch.tensor(lam_batch).unsqueeze(1)
+
+    def _mix_pair_collate(self, output, batch):
+        batch_size = len(batch)
+        lam_batch, use_cutmix = self._params_per_elem(batch_size // 2)
+        for i in range(batch_size // 2):
+            j = batch_size - i - 1
+            lam = lam_batch[i]
+            mixed_i = batch[i][0]
+            mixed_j = batch[j][0]
+            assert 0 <= lam <= 1.0
+            if lam < 1.:
+                if use_cutmix[i]:
+                    (yl, yh, xl, xh), lam = cutmix_bbox_and_lam(
+                        output.shape, lam, ratio_minmax=self.cutmix_minmax, correct_lam=self.correct_lam)
+                    patch_i = mixed_i[:, yl:yh, xl:xh].copy()
+                    mixed_i[:, yl:yh, xl:xh] = mixed_j[:, yl:yh, xl:xh]
+                    mixed_j[:, yl:yh, xl:xh] = patch_i
+                    lam_batch[i] = lam
+                else:
+                    mixed_temp = mixed_i.astype(np.float32) * lam + mixed_j.astype(np.float32) * (1 - lam)
+                    mixed_j = mixed_j.astype(np.float32) * lam + mixed_i.astype(np.float32) * (1 - lam)
+                    mixed_i = mixed_temp
+                    np.rint(mixed_j, out=mixed_j)
+                    np.rint(mixed_i, out=mixed_i)
+            output[i] += torch.from_numpy(mixed_i.astype(np.uint8))
+            output[j] += torch.from_numpy(mixed_j.astype(np.uint8))
+        lam_batch = np.concatenate((lam_batch, lam_batch[::-1]))
+        return torch.tensor(lam_batch).unsqueeze(1)
+
+    def _mix_batch_collate(self, output, batch):
+        batch_size = len(batch)
+        lam, use_cutmix = self._params_per_batch()
+        if use_cutmix:
+            (yl, yh, xl, xh), lam = cutmix_bbox_and_lam(
+                output.shape, lam, ratio_minmax=self.cutmix_minmax, correct_lam=self.correct_lam)
+        for i in range(batch_size):
+            j = batch_size - i - 1
+            mixed = batch[i][0]
+            if lam != 1.:
+                if use_cutmix:
+                    mixed = mixed.copy()  # don't want to modify the original while iterating
+                    mixed[:, yl:yh, xl:xh] = batch[j][0][:, yl:yh, xl:xh]
+                else:
+                    mixed = mixed.astype(np.float32) * lam + batch[j][0].astype(np.float32) * (1 - lam)
+                    np.rint(mixed, out=mixed)
+            output[i] += torch.from_numpy(mixed.astype(np.uint8))
+        return lam
+
+    def __call__(self, batch, _=None):
+        batch_size = len(batch)
+        assert batch_size % 2 == 0, 'Batch size should be even when using this'
+        half = 'half' in self.mode
+        if half:
+            batch_size //= 2
+        output = torch.zeros((batch_size, *batch[0][0].shape), dtype=torch.uint8)
+        if self.mode == 'elem' or self.mode == 'half':
+            lam = self._mix_elem_collate(output, batch, half=half)
+        elif self.mode == 'pair':
+            lam = self._mix_pair_collate(output, batch)
+        else:
+            lam = self._mix_batch_collate(output, batch)
+        target = torch.tensor([b[1] for b in batch], dtype=torch.int64)
+        target = mixup_target(target, self.num_classes, lam, self.label_smoothing)
+        target = target[:batch_size]
+        return output, target

timm/data/random_erasing.py

@@ -0,0 +1,117 @@
+""" Random Erasing (Cutout)
+
+Originally inspired by impl at https://github.com/zhunzhong07/Random-Erasing, Apache 2.0
+Copyright Zhun Zhong & Liang Zheng
+
+Hacked together by / Copyright 2019, Ross Wightman
+"""
+import random
+import math
+
+import torch
+
+
+def _get_pixels(per_pixel, rand_color, patch_size, dtype=torch.float32, device='cuda'):
+    # NOTE I've seen CUDA illegal memory access errors being caused by the normal_()
+    # paths, flip the order so normal is run on CPU if this becomes a problem
+    # Issue has been fixed in master https://github.com/pytorch/pytorch/issues/19508
+    if per_pixel:
+        return torch.empty(patch_size, dtype=dtype, device=device).normal_()
+    elif rand_color:
+        return torch.empty((patch_size[0], 1, 1), dtype=dtype, device=device).normal_()
+    else:
+        return torch.zeros((patch_size[0], 1, 1), dtype=dtype, device=device)
+
+
+class RandomErasing:
+    """ Randomly selects a rectangle region in an image and erases its pixels.
+        'Random Erasing Data Augmentation' by Zhong et al.
+        See https://arxiv.org/pdf/1708.04896.pdf
+
+        This variant of RandomErasing is intended to be applied to either a batch
+        or single image tensor after it has been normalized by dataset mean and std.
+    Args:
+         probability: Probability that the Random Erasing operation will be performed.
+         min_area: Minimum percentage of erased area wrt input image area.
+         max_area: Maximum percentage of erased area wrt input image area.
+         min_aspect: Minimum aspect ratio of erased area.
+         mode: pixel color mode, one of 'const', 'rand', or 'pixel'
+            'const' - erase block is constant color of 0 for all channels
+            'rand'  - erase block is same per-channel random (normal) color
+            'pixel' - erase block is per-pixel random (normal) color
+        max_count: maximum number of erasing blocks per image, area per box is scaled by count.
+            per-image count is randomly chosen between 1 and this value.
+    """
+
+    def __init__(
+            self,
+            probability=0.5,
+            min_area=0.02,
+            max_area=1/3,
+            min_aspect=0.3,
+            max_aspect=None,
+            mode='const',
+            min_count=1,
+            max_count=None,
+            num_splits=0,
+            device='cuda',
+    ):
+        self.probability = probability
+        self.min_area = min_area
+        self.max_area = max_area
+        max_aspect = max_aspect or 1 / min_aspect
+        self.log_aspect_ratio = (math.log(min_aspect), math.log(max_aspect))
+        self.min_count = min_count
+        self.max_count = max_count or min_count
+        self.num_splits = num_splits
+        self.mode = mode.lower()
+        self.rand_color = False
+        self.per_pixel = False
+        if self.mode == 'rand':
+            self.rand_color = True  # per block random normal
+        elif self.mode == 'pixel':
+            self.per_pixel = True  # per pixel random normal
+        else:
+            assert not self.mode or self.mode == 'const'
+        self.device = device
+
+    def _erase(self, img, chan, img_h, img_w, dtype):
+        if random.random() > self.probability:
+            return
+        area = img_h * img_w
+        count = self.min_count if self.min_count == self.max_count else \
+            random.randint(self.min_count, self.max_count)
+        for _ in range(count):
+            for attempt in range(10):
+                target_area = random.uniform(self.min_area, self.max_area) * area / count
+                aspect_ratio = math.exp(random.uniform(*self.log_aspect_ratio))
+                h = int(round(math.sqrt(target_area * aspect_ratio)))
+                w = int(round(math.sqrt(target_area / aspect_ratio)))
+                if w < img_w and h < img_h:
+                    top = random.randint(0, img_h - h)
+                    left = random.randint(0, img_w - w)
+                    img[:, top:top + h, left:left + w] = _get_pixels(
+                        self.per_pixel,
+                        self.rand_color,
+                        (chan, h, w),
+                        dtype=dtype,
+                        device=self.device,
+                    )
+                    break
+
+    def __call__(self, input):
+        if len(input.size()) == 3:
+            self._erase(input, *input.size(), input.dtype)
+        else:
+            batch_size, chan, img_h, img_w = input.size()
+            # skip first slice of batch if num_splits is set (for clean portion of samples)
+            batch_start = batch_size // self.num_splits if self.num_splits > 1 else 0
+            for i in range(batch_start, batch_size):
+                self._erase(input[i], chan, img_h, img_w, input.dtype)
+        return input
+
+    def __repr__(self):
+        # NOTE simplified state for repr
+        fs = self.__class__.__name__ + f'(p={self.probability}, mode={self.mode}'
+        fs += f', count=({self.min_count}, {self.max_count}))'
+        return fs

timm/data/readers/__init__.py

@@ -0,0 +1,2 @@
+from .reader_factory import create_reader
+from .img_extensions import *

timm/data/readers/class_map.py

@@ -0,0 +1,22 @@
+import os
+import pickle
+
+
+def load_class_map(map_or_filename, root=''):
+    if isinstance(map_or_filename, dict):
+        assert dict, 'class_map dict must be non-empty'
+        return map_or_filename
+    class_map_path = map_or_filename
+    if not os.path.exists(class_map_path):
+        class_map_path = os.path.join(root, class_map_path)
+        assert os.path.exists(class_map_path), 'Cannot locate specified class map file (%s)' % map_or_filename
+    class_map_ext = os.path.splitext(map_or_filename)[-1].lower()
+    if class_map_ext == '.txt':
+        with open(class_map_path) as f:
+            class_to_idx = {v.strip(): k for k, v in enumerate(f)}
+    elif class_map_ext == '.pkl':
+        with open(class_map_path, 'rb') as f:
+            class_to_idx = pickle.load(f)
+    else:
+        assert False, f'Unsupported class map file extension ({class_map_ext}).'
+    return class_to_idx

timm/data/readers/img_extensions.py

@@ -0,0 +1,50 @@
+from copy import deepcopy
+
+__all__ = ['get_img_extensions', 'is_img_extension', 'set_img_extensions', 'add_img_extensions', 'del_img_extensions']
+
+
+IMG_EXTENSIONS = ('.png', '.jpg', '.jpeg')  # singleton, kept public for bwd compat use
+_IMG_EXTENSIONS_SET = set(IMG_EXTENSIONS)  # set version, private, kept in sync
+
+
+def _set_extensions(extensions):
+    global IMG_EXTENSIONS
+    global _IMG_EXTENSIONS_SET
+    dedupe = set()  # NOTE de-duping tuple while keeping original order
+    IMG_EXTENSIONS = tuple(x for x in extensions if x not in dedupe and not dedupe.add(x))
+    _IMG_EXTENSIONS_SET = set(extensions)
+
+
+def _valid_extension(x: str):
+    return x and isinstance(x, str) and len(x) >= 2 and x.startswith('.')
+
+
+def is_img_extension(ext):
+    return ext in _IMG_EXTENSIONS_SET
+
+
+def get_img_extensions(as_set=False):
+    return deepcopy(_IMG_EXTENSIONS_SET if as_set else IMG_EXTENSIONS)
+
+
+def set_img_extensions(extensions):
+    assert len(extensions)
+    for x in extensions:
+        assert _valid_extension(x)
+    _set_extensions(extensions)
+
+
+def add_img_extensions(ext):
+    if not isinstance(ext, (list, tuple, set)):
+        ext = (ext,)
+    for x in ext:
+        assert _valid_extension(x)
+    extensions = IMG_EXTENSIONS + tuple(ext)
+    _set_extensions(extensions)
+
+
+def del_img_extensions(ext):
+    if not isinstance(ext, (list, tuple, set)):
+        ext = (ext,)
+    extensions = tuple(x for x in IMG_EXTENSIONS if x not in ext)
+    _set_extensions(extensions)

timm/data/readers/reader.py

@@ -0,0 +1,16 @@
+from abc import abstractmethod
+
+
+class Reader:
+    def __init__(self):
+        pass
+
+    @abstractmethod
+    def _filename(self, index, basename=False, absolute=False):
+        pass
+
+    def filename(self, index, basename=False, absolute=False):
+        return self._filename(index, basename=basename, absolute=absolute)
+
+    def filenames(self, basename=False, absolute=False):
+        return [self._filename(index, basename=basename, absolute=absolute) for index in range(len(self))]

timm/data/readers/reader_factory.py

@@ -0,0 +1,45 @@
+import os
+from typing import Optional
+
+from .reader_image_folder import ReaderImageFolder
+from .reader_image_in_tar import ReaderImageInTar
+
+
+def create_reader(
+        name: str,
+        root: Optional[str] = None,
+        split: str = 'train',
+        **kwargs,
+):
+    kwargs = {k: v for k, v in kwargs.items() if v is not None}
+    name = name.lower()
+    name = name.split('/', 1)
+    prefix = ''
+    if len(name) > 1:
+        prefix = name[0]
+    name = name[-1]
+
+    # FIXME improve the selection right now just tfds prefix or fallback path, will need options to
+    # explicitly select other options shortly
+    if prefix == 'hfds':
+        from .reader_hfds import ReaderHfds  # defer Hf datasets import
+        reader = ReaderHfds(name=name, root=root, split=split, **kwargs)
+    elif prefix == 'hfids':
+        from .reader_hfids import ReaderHfids  # defer HF datasets import
+        reader = ReaderHfids(name=name, root=root, split=split, **kwargs)
+    elif prefix == 'tfds':
+        from .reader_tfds import ReaderTfds  # defer tensorflow import
+        reader = ReaderTfds(name=name, root=root, split=split, **kwargs)
+    elif prefix == 'wds':
+        from .reader_wds import ReaderWds
+        kwargs.pop('download', False)
+        reader = ReaderWds(root=root, name=name, split=split, **kwargs)
+    else:
+        assert os.path.exists(root)
+        # default fallback path (backwards compat), use image tar if root is a .tar file, otherwise image folder
+        # FIXME support split here or in reader?
+        if os.path.isfile(root) and os.path.splitext(root)[1] == '.tar':
+            reader = ReaderImageInTar(root, **kwargs)
+        else:
+            reader = ReaderImageFolder(root, **kwargs)
+    return reader

timm/data/readers/reader_hfds.py

@@ -0,0 +1,84 @@
+""" Dataset reader that wraps Hugging Face datasets
+
+Hacked together by / Copyright 2022 Ross Wightman
+"""
+import io
+import math
+from typing import Optional
+
+import torch
+import torch.distributed as dist
+from PIL import Image
+
+try:
+    import datasets
+except ImportError as e:
+    print("Please install Hugging Face datasets package `pip install datasets`.")
+    raise e
+from .class_map import load_class_map
+from .reader import Reader
+
+
+def get_class_labels(info, label_key='label'):
+    if 'label' not in info.features:
+        return {}
+    class_label = info.features[label_key]
+    class_to_idx = {n: class_label.str2int(n) for n in class_label.names}
+    return class_to_idx
+
+
+class ReaderHfds(Reader):
+
+    def __init__(
+            self,
+            name: str,
+            root: Optional[str] = None,
+            split: str = 'train',
+            class_map: dict = None,
+            image_key: str = 'image',
+            target_key: str = 'label',
+            download: bool = False,
+    ):
+        """
+        """
+        super().__init__()
+        self.root = root
+        self.split = split
+        self.dataset = datasets.load_dataset(
+            name,  # 'name' maps to path arg in hf datasets
+            split=split,
+            cache_dir=self.root,  # timm doesn't expect hidden cache dir for datasets, specify a path
+        )
+        # leave decode for caller, plus we want easy access to original path names...
+        self.dataset = self.dataset.cast_column(image_key, datasets.Image(decode=False))
+
+        self.image_key = image_key
+        self.label_key = target_key
+        self.remap_class = False
+        if class_map:
+            self.class_to_idx = load_class_map(class_map)
+            self.remap_class = True
+        else:
+            self.class_to_idx = get_class_labels(self.dataset.info, self.label_key)
+        self.split_info = self.dataset.info.splits[split]
+        self.num_samples = self.split_info.num_examples
+
+    def __getitem__(self, index):
+        item = self.dataset[index]
+        image = item[self.image_key]
+        if 'bytes' in image and image['bytes']:
+            image = io.BytesIO(image['bytes'])
+        else:
+            assert 'path' in image and image['path']
+            image = open(image['path'], 'rb')
+        label = item[self.label_key]
+        if self.remap_class:
+            label = self.class_to_idx[label]
+        return image, label
+
+    def __len__(self):
+        return len(self.dataset)
+
+    def _filename(self, index, basename=False, absolute=False):
+        item = self.dataset[index]
+        return item[self.image_key]['path']

timm/data/readers/reader_hfids.py

@@ -0,0 +1,215 @@
+""" Dataset reader for HF IterableDataset
+"""
+import math
+import os
+from itertools import repeat, chain
+from typing import Optional
+
+import torch
+import torch.distributed as dist
+from PIL import Image
+
+try:
+    import datasets
+    from datasets.distributed import split_dataset_by_node
+    from datasets.splits import SplitInfo
+except ImportError as e:
+    print("Please install Hugging Face datasets package `pip install datasets`.")
+    raise e
+
+
+from .class_map import load_class_map
+from .reader import Reader
+from .shared_count import SharedCount
+
+
+SHUFFLE_SIZE = int(os.environ.get('HFIDS_SHUFFLE_SIZE', 4096))
+
+
+class ReaderHfids(Reader):
+    def __init__(
+            self,
+            name: str,
+            root: Optional[str] = None,
+            split: str = 'train',
+            is_training: bool = False,
+            batch_size: int = 1,
+            download: bool = False,
+            repeats: int = 0,
+            seed: int = 42,
+            class_map: Optional[dict] = None,
+            input_key: str = 'image',
+            input_img_mode: str = 'RGB',
+            target_key: str = 'label',
+            target_img_mode: str = '',
+            shuffle_size: Optional[int] = None,
+            num_samples: Optional[int] = None,
+    ):
+        super().__init__()
+        self.root = root
+        self.split = split
+        self.is_training = is_training
+        self.batch_size = batch_size
+        self.download = download
+        self.repeats = repeats
+        self.common_seed = seed  # a seed that's fixed across all worker / distributed instances
+        self.shuffle_size = shuffle_size or SHUFFLE_SIZE
+
+        self.input_key = input_key
+        self.input_img_mode = input_img_mode
+        self.target_key = target_key
+        self.target_img_mode = target_img_mode
+
+        self.builder = datasets.load_dataset_builder(name, cache_dir=root)
+        if download:
+            self.builder.download_and_prepare()
+
+        split_info: Optional[SplitInfo] = None
+        if self.builder.info.splits and split in self.builder.info.splits:
+            if isinstance(self.builder.info.splits[split], SplitInfo):
+                split_info: Optional[SplitInfo] = self.builder.info.splits[split]
+
+        if num_samples:
+            self.num_samples = num_samples
+        elif split_info and split_info.num_examples:
+            self.num_samples = split_info.num_examples
+        else:
+            raise ValueError(
+                "Dataset length is unknown, please pass `num_samples` explicitely. "
+                "The number of steps needs to be known in advance for the learning rate scheduler."
+            )
+
+        self.remap_class = False
+        if class_map:
+            self.class_to_idx = load_class_map(class_map)
+            self.remap_class = True
+        else:
+            self.class_to_idx = {}
+
+        # Distributed world state
+        self.dist_rank = 0
+        self.dist_num_replicas = 1
+        if dist.is_available() and dist.is_initialized() and dist.get_world_size() > 1:
+            self.dist_rank = dist.get_rank()
+            self.dist_num_replicas = dist.get_world_size()
+
+        # Attributes that are updated in _lazy_init
+        self.worker_info = None
+        self.worker_id = 0
+        self.num_workers = 1
+        self.global_worker_id = 0
+        self.global_num_workers = 1
+
+        # Initialized lazily on each dataloader worker process
+        self.ds: Optional[datasets.IterableDataset] = None
+        self.epoch = SharedCount()
+
+    def set_epoch(self, count):
+        # to update the shuffling effective_seed = seed + epoch
+        self.epoch.value = count
+
+    def set_loader_cfg(
+            self,
+            num_workers: Optional[int] = None,
+    ):
+        if self.ds is not None:
+            return
+        if num_workers is not None:
+            self.num_workers = num_workers
+            self.global_num_workers = self.dist_num_replicas * self.num_workers
+
+    def _lazy_init(self):
+        """ Lazily initialize worker (in worker processes)
+        """
+        if self.worker_info is None:
+            worker_info = torch.utils.data.get_worker_info()
+            if worker_info is not None:
+                self.worker_info = worker_info
+                self.worker_id = worker_info.id
+                self.num_workers = worker_info.num_workers
+            self.global_num_workers = self.dist_num_replicas * self.num_workers
+            self.global_worker_id = self.dist_rank * self.num_workers + self.worker_id
+
+        if self.download:
+            dataset = self.builder.as_dataset(split=self.split)
+            # to distribute evenly to workers
+            ds = dataset.to_iterable_dataset(num_shards=self.global_num_workers)
+        else:
+            # in this case the number of shard is determined by the number of remote files
+            ds = self.builder.as_streaming_dataset(split=self.split)
+
+        if self.is_training:
+            # will shuffle the list of shards and use a shuffle buffer
+            ds = ds.shuffle(seed=self.common_seed, buffer_size=self.shuffle_size)
+
+        # Distributed:
+        # The dataset has a number of shards that is a factor of `dist_num_replicas` (i.e. if `ds.n_shards % dist_num_replicas == 0`),
+        # so the shards are evenly assigned across the nodes.
+        # If it's not the case for dataset streaming, each node keeps 1 example out of `dist_num_replicas`, skipping the other examples.
+
+        # Workers:
+        # In a node, datasets.IterableDataset assigns the shards assigned to the node as evenly as possible to workers.
+        self.ds = split_dataset_by_node(ds, rank=self.dist_rank, world_size=self.dist_num_replicas)
+
+    def _num_samples_per_worker(self):
+        num_worker_samples = \
+            max(1, self.repeats) * self.num_samples / max(self.global_num_workers, self.dist_num_replicas)
+        if self.is_training or self.dist_num_replicas > 1:
+            num_worker_samples = math.ceil(num_worker_samples)
+        if self.is_training and self.batch_size is not None:
+            num_worker_samples = math.ceil(num_worker_samples / self.batch_size) * self.batch_size
+        return int(num_worker_samples)
+
+    def __iter__(self):
+        if self.ds is None:
+            self._lazy_init()
+        self.ds.set_epoch(self.epoch.value)
+
+        target_sample_count = self._num_samples_per_worker()
+        sample_count = 0
+
+        if self.is_training:
+            ds_iter = chain.from_iterable(repeat(self.ds))
+        else:
+            ds_iter = iter(self.ds)
+        for sample in ds_iter:
+            input_data: Image.Image = sample[self.input_key]
+            if self.input_img_mode and input_data.mode != self.input_img_mode:
+                input_data = input_data.convert(self.input_img_mode)
+            target_data = sample[self.target_key]
+            if self.target_img_mode:
+                assert isinstance(target_data, Image.Image), "target_img_mode is specified but target is not an image"
+                if target_data.mode != self.target_img_mode:
+                    target_data = target_data.convert(self.target_img_mode)
+            elif self.remap_class:
+                target_data = self.class_to_idx[target_data]
+            yield input_data, target_data
+            sample_count += 1
+            if self.is_training and sample_count >= target_sample_count:
+                break
+
+    def __len__(self):
+        num_samples = self._num_samples_per_worker() * self.num_workers
+        return num_samples
+
+    def _filename(self, index, basename=False, absolute=False):
+        assert False, "Not supported"  # no random access to examples
+
+    def filenames(self, basename=False, absolute=False):
+        """ Return all filenames in dataset, overrides base"""
+        if self.ds is None:
+            self._lazy_init()
+        names = []
+        for sample in self.ds:
+            if 'file_name' in sample:
+                name = sample['file_name']
+            elif 'filename' in sample:
+                name = sample['filename']
+            elif 'id' in sample:
+                name = sample['id']
+            elif 'image_id' in sample:
+                name = sample['image_id']
+            else:
+                assert False, "No supported name field present"
+            names.append(name)
+        return names

timm/data/readers/reader_image_folder.py

@@ -0,0 +1,99 @@
+""" A dataset reader that extracts images from folders
+
+Folders are scanned recursively to find image files. Labels are based
+on the folder hierarchy, just leaf folders by default.
+
+Hacked together by / Copyright 2020 Ross Wightman
+"""
+import os
+from typing import Dict, List, Optional, Set, Tuple, Union
+
+from timm.utils.misc import natural_key
+
+from .class_map import load_class_map
+from .img_extensions import get_img_extensions
+from .reader import Reader
+
+
+def find_images_and_targets(
+        folder: str,
+        types: Optional[Union[List, Tuple, Set]] = None,
+        class_to_idx: Optional[Dict] = None,
+        leaf_name_only: bool = True,
+        sort: bool = True
+):
+    """ Walk folder recursively to discover images and map them to classes by folder names.
+
+    Args:
+        folder: root of folder to recrusively search
+        types: types (file extensions) to search for in path
+        class_to_idx: specify mapping for class (folder name) to class index if set
+        leaf_name_only: use only leaf-name of folder walk for class names
+        sort: re-sort found images by name (for consistent ordering)
+
+    Returns:
+        A list of image and target tuples, class_to_idx mapping
+    """
+    types = get_img_extensions(as_set=True) if not types else set(types)
+    labels = []
+    filenames = []
+    for root, subdirs, files in os.walk(folder, topdown=False, followlinks=True):
+        rel_path = os.path.relpath(root, folder) if (root != folder) else ''
+        label = os.path.basename(rel_path) if leaf_name_only else rel_path.replace(os.path.sep, '_')
+        for f in files:
+            base, ext = os.path.splitext(f)
+            if ext.lower() in types:
+                filenames.append(os.path.join(root, f))
+                labels.append(label)
+    if class_to_idx is None:
+        # building class index
+        unique_labels = set(labels)
+        sorted_labels = list(sorted(unique_labels, key=natural_key))
+        class_to_idx = {c: idx for idx, c in enumerate(sorted_labels)}
+    images_and_targets = [(f, class_to_idx[l]) for f, l in zip(filenames, labels) if l in class_to_idx]
+    if sort:
+        images_and_targets = sorted(images_and_targets, key=lambda k: natural_key(k[0]))
+    return images_and_targets, class_to_idx
+
+
+class ReaderImageFolder(Reader):
+
+    def __init__(
+            self,
+            root,
+            class_map='',
+            input_key=None,
+    ):
+        super().__init__()
+
+        self.root = root
+        class_to_idx = None
+        if class_map:
+            class_to_idx = load_class_map(class_map, root)
+        find_types = None
+        if input_key:
+            find_types = input_key.split(';')
+        self.samples, self.class_to_idx = find_images_and_targets(
+            root,
+            class_to_idx=class_to_idx,
+            types=find_types,
+        )
+        if len(self.samples) == 0:
+            raise RuntimeError(
+                f'Found 0 images in subfolders of {root}. '
+                f'Supported image extensions are {", ".join(get_img_extensions())}')
+
+    def __getitem__(self, index):
+        path, target = self.samples[index]
+        return open(path, 'rb'), target
+
+    def __len__(self):
+        return len(self.samples)
+
+    def _filename(self, index, basename=False, absolute=False):
+        filename = self.samples[index][0]
+        if basename:
+            filename = os.path.basename(filename)
+        elif not absolute:
+            filename = os.path.relpath(filename, self.root)
+        return filename

timm/data/readers/reader_image_in_tar.py

@@ -0,0 +1,229 @@
+""" A dataset reader that reads tarfile based datasets
+
+This reader can extract image samples from:
+* a single tar of image files
+* a folder of multiple tarfiles containing imagefiles
+* a tar of tars containing image files
+
+Labels are based on the combined folder and/or tar name structure.
+
+Hacked together by / Copyright 2020 Ross Wightman
+"""
+import logging
+import os
+import pickle
+import tarfile
+from glob import glob
+from typing import List, Tuple, Dict, Set, Optional, Union
+
+import numpy as np
+
+from timm.utils.misc import natural_key
+
+from .class_map import load_class_map
+from .img_extensions import get_img_extensions
+from .reader import Reader
+
+_logger = logging.getLogger(__name__)
+CACHE_FILENAME_SUFFIX = '_tarinfos.pickle'
+
+
+class TarState:
+
+    def __init__(self, tf: tarfile.TarFile = None, ti: tarfile.TarInfo = None):
+        self.tf: tarfile.TarFile = tf
+        self.ti: tarfile.TarInfo = ti
+        self.children: Dict[str, TarState] = {}  # child states (tars within tars)
+
+    def reset(self):
+        self.tf = None
+
+
+def _extract_tarinfo(tf: tarfile.TarFile, parent_info: Dict, extensions: Set[str]):
+    sample_count = 0
+    for i, ti in enumerate(tf):
+        if not ti.isfile():
+            continue
+        dirname, basename = os.path.split(ti.path)
+        name, ext = os.path.splitext(basename)
+        ext = ext.lower()
+        if ext == '.tar':
+            with tarfile.open(fileobj=tf.extractfile(ti), mode='r|') as ctf:
+                child_info = dict(
+                    name=ti.name, path=os.path.join(parent_info['path'], name), ti=ti, children=[], samples=[])
+                sample_count += _extract_tarinfo(ctf, child_info, extensions=extensions)
+                _logger.debug(f'{i}/?. Extracted child tarinfos from {ti.name}. {len(child_info["samples"])} images.')
+                parent_info['children'].append(child_info)
+        elif ext in extensions:
+            parent_info['samples'].append(ti)
+            sample_count += 1
+    return sample_count
+
+
+def extract_tarinfos(
+        root,
+        class_name_to_idx: Optional[Dict] = None,
+        cache_tarinfo: Optional[bool] = None,
+        extensions: Optional[Union[List, Tuple, Set]] = None,
+        sort: bool = True
+):
+    extensions = get_img_extensions(as_set=True) if not extensions else set(extensions)
+    root_is_tar = False
+    if os.path.isfile(root):
+        assert os.path.splitext(root)[-1].lower() == '.tar'
+        tar_filenames = [root]
+        root, root_name = os.path.split(root)
+        root_name = os.path.splitext(root_name)[0]
+        root_is_tar = True
+    else:
+        root_name = root.strip(os.path.sep).split(os.path.sep)[-1]
+        tar_filenames = glob(os.path.join(root, '*.tar'), recursive=True)
+    num_tars = len(tar_filenames)
+    tar_bytes = sum([os.path.getsize(f) for f in tar_filenames])
+    assert num_tars, f'No .tar files found at specified path ({root}).'
+
+    _logger.info(f'Scanning {tar_bytes/1024**2:.2f}MB of tar files...')
+    info = dict(tartrees=[])
+    cache_path = ''
+    if cache_tarinfo is None:
+        cache_tarinfo = True if tar_bytes > 10*1024**3 else False  # FIXME magic number, 10GB
+    if cache_tarinfo:
+        cache_filename = '_' + root_name + CACHE_FILENAME_SUFFIX
+        cache_path = os.path.join(root, cache_filename)
+    if os.path.exists(cache_path):
+        _logger.info(f'Reading tar info from cache file {cache_path}.')
+        with open(cache_path, 'rb') as pf:
+            info = pickle.load(pf)
+        assert len(info['tartrees']) == num_tars, "Cached tartree len doesn't match number of tarfiles"
+    else:
+        for i, fn in enumerate(tar_filenames):
+            path = '' if root_is_tar else os.path.splitext(os.path.basename(fn))[0]
+            with tarfile.open(fn, mode='r|') as tf:  # tarinfo scans done in streaming mode
+                parent_info = dict(name=os.path.relpath(fn, root), path=path, ti=None, children=[], samples=[])
+                num_samples = _extract_tarinfo(tf, parent_info, extensions=extensions)
+                num_children = len(parent_info["children"])
+                _logger.debug(
+                    f'{i}/{num_tars}. Extracted tarinfos from {fn}. {num_children} children, {num_samples} samples.')
+            info['tartrees'].append(parent_info)
+        if cache_path:
+            _logger.info(f'Writing tar info to cache file {cache_path}.')
+            with open(cache_path, 'wb') as pf:
+                pickle.dump(info, pf)
+
+    samples = []
+    labels = []
+    build_class_map = False
+    if class_name_to_idx is None:
+        build_class_map = True
+
+    # Flatten tartree info into lists of samples and targets w/ targets based on label id via
+    # class map arg or from unique paths.
+    # NOTE: currently only flattening up to two-levels, filesystem .tars and then one level of sub-tar children
+    # this covers my current use cases and keeps things a little easier to test for now.
+    tarfiles = []
+
+    def _label_from_paths(*path, leaf_only=True):
+        path = os.path.join(*path).strip(os.path.sep)
+        return path.split(os.path.sep)[-1] if leaf_only else path.replace(os.path.sep, '_')
+
+    def _add_samples(info, fn):
+        added = 0
+        for s in info['samples']:
+            label = _label_from_paths(info['path'], os.path.dirname(s.path))
+            if not build_class_map and label not in class_name_to_idx:
+                continue
+            samples.append((s, fn, info['ti']))
+            labels.append(label)
+            added += 1
+        return added
+
+    _logger.info(f'Collecting samples and building tar states.')
+    for parent_info in info['tartrees']:
+        # if tartree has children, we assume all samples are at the child level
+        tar_name = None if root_is_tar else parent_info['name']
+        tar_state = TarState()
+        parent_added = 0
+        for child_info in parent_info['children']:
+            child_added = _add_samples(child_info, fn=tar_name)
+            if child_added:
+                tar_state.children[child_info['name']] = TarState(ti=child_info['ti'])
+            parent_added += child_added
+        parent_added += _add_samples(parent_info, fn=tar_name)
+        if parent_added:
+            tarfiles.append((tar_name, tar_state))
+    del info
+
+    if build_class_map:
+        # build class index
+        sorted_labels = list(sorted(set(labels), key=natural_key))
+        class_name_to_idx = {c: idx for idx, c in enumerate(sorted_labels)}
+
+    _logger.info(f'Mapping targets and sorting samples.')
+    samples_and_targets = [(s, class_name_to_idx[l]) for s, l in zip(samples, labels) if l in class_name_to_idx]
+    if sort:
+        samples_and_targets = sorted(samples_and_targets, key=lambda k: natural_key(k[0][0].path))
+    samples, targets = zip(*samples_and_targets)
+    samples = np.array(samples)
+    targets = np.array(targets)
+    _logger.info(f'Finished processing {len(samples)} samples across {len(tarfiles)} tar files.')
+    return samples, targets, class_name_to_idx, tarfiles
+
+
+class ReaderImageInTar(Reader):
+    """ Multi-tarfile dataset reader where there is one .tar file per class
+    """
+
+    def __init__(self, root, class_map='', cache_tarfiles=True, cache_tarinfo=None):
+        super().__init__()
+
+        class_name_to_idx = None
+        if class_map:
+            class_name_to_idx = load_class_map(class_map, root)
+        self.root = root
+        self.samples, self.targets, self.class_name_to_idx, tarfiles = extract_tarinfos(
+            self.root,
+            class_name_to_idx=class_name_to_idx,
+            cache_tarinfo=cache_tarinfo
+        )
+        self.class_idx_to_name = {v: k for k, v in self.class_name_to_idx.items()}
+        if len(tarfiles) == 1 and tarfiles[0][0] is None:
+            self.root_is_tar = True
+            self.tar_state = tarfiles[0][1]
+        else:
+            self.root_is_tar = False
+            self.tar_state = dict(tarfiles)
+        self.cache_tarfiles = cache_tarfiles
+
+    def __len__(self):
+        return len(self.samples)
+
+    def __getitem__(self, index):
+        sample = self.samples[index]
+        target = self.targets[index]
+        sample_ti, parent_fn, child_ti = sample
+        parent_abs = os.path.join(self.root, parent_fn) if parent_fn else self.root
+
+        tf = None
+        cache_state = None
+        if self.cache_tarfiles:
+            cache_state = self.tar_state if self.root_is_tar else self.tar_state[parent_fn]
+            tf = cache_state.tf
+        if tf is None:
+            tf = tarfile.open(parent_abs)
+            if self.cache_tarfiles:
+                cache_state.tf = tf
+        if child_ti is not None:
+            ctf = cache_state.children[child_ti.name].tf if self.cache_tarfiles else None
+            if ctf is None:
+                ctf = tarfile.open(fileobj=tf.extractfile(child_ti))
+                if self.cache_tarfiles:
+                    cache_state.children[child_ti.name].tf = ctf
+            tf = ctf
+
+        return tf.extractfile(sample_ti), target
+
+    def _filename(self, index, basename=False, absolute=False):
+        filename = self.samples[index][0].name
+        if basename:
+            filename = os.path.basename(filename)
+        return filename

timm/data/readers/reader_tfds.py

@@ -0,0 +1,355 @@
+""" Dataset reader that wraps TFDS datasets
+
+Wraps many (most?) TFDS image-classification datasets
+from https://github.com/tensorflow/datasets
+https://www.tensorflow.org/datasets/catalog/overview#image_classification
+
+Hacked together by / Copyright 2020 Ross Wightman
+"""
+import math
+import os
+import sys
+from typing import Optional
+
+import torch
+import torch.distributed as dist
+from PIL import Image
+
+try:
+    import tensorflow as tf
+    tf.config.set_visible_devices([], 'GPU')  # Hands off my GPU! (or pip install tensorflow-cpu)
+    import tensorflow_datasets as tfds
+    try:
+        tfds.even_splits('', 1, drop_remainder=False)  # non-buggy even_splits has drop_remainder arg
+        has_buggy_even_splits = False
+    except TypeError:
+        print("Warning: This version of tfds doesn't have the latest even_splits impl. "
+              "Please update or use tfds-nightly for better fine-grained split behaviour.")
+        has_buggy_even_splits = True
+    # NOTE uncomment below if having file limit issues on dataset build (or alter your OS defaults)
+    # import resource
+    # low, high = resource.getrlimit(resource.RLIMIT_NOFILE)
+    # resource.setrlimit(resource.RLIMIT_NOFILE, (high, high))
+except ImportError as e:
+    print(e)
+    print("Please install tensorflow_datasets package `pip install tensorflow-datasets`.")
+    raise e
+
+from .class_map import load_class_map
+from .reader import Reader
+from .shared_count import SharedCount
+
+
+MAX_TP_SIZE = int(os.environ.get('TFDS_TP_SIZE', 8))  # maximum TF threadpool size, for jpeg decodes and queuing activities
+SHUFFLE_SIZE = int(os.environ.get('TFDS_SHUFFLE_SIZE', 8192))  # samples to shuffle in DS queue
+PREFETCH_SIZE = int(os.environ.get('TFDS_PREFETCH_SIZE', 2048))  # samples to prefetch
+
+
+@tfds.decode.make_decoder()
+def decode_example(serialized_image, feature, dct_method='INTEGER_ACCURATE', channels=3):
+    return tf.image.decode_jpeg(
+        serialized_image,
+        channels=channels,
+        dct_method=dct_method,
+    )
+
+
+def even_split_indices(split, n, num_samples):
+    partitions = [round(i * num_samples / n) for i in range(n + 1)]
+    return [f"{split}[{partitions[i]}:{partitions[i + 1]}]" for i in range(n)]
+
+
+def get_class_labels(info):
+    if 'label' not in info.features:
+        return {}
+    class_label = info.features['label']
+    class_to_idx = {n: class_label.str2int(n) for n in class_label.names}
+    return class_to_idx
+
+
+class ReaderTfds(Reader):
+    """ Wrap Tensorflow Datasets for use in PyTorch
+
+    There several things to be aware of:
+      * To prevent excessive samples being dropped per epoch w/ distributed training or multiplicity of
+         dataloader workers, the train iterator wraps to avoid returning partial batches that trigger drop_last
+         https://github.com/pytorch/pytorch/issues/33413
+      * With PyTorch IterableDatasets, each worker in each replica operates in isolation, the final batch
+        from each worker could be a different size. For training this is worked around by option above, for
+        validation extra samples are inserted iff distributed mode is enabled so that the batches being reduced
+        across replicas are of same size. This will slightly alter the results, distributed validation will not be
+        100% correct. This is similar to common handling in DistributedSampler for normal Datasets but a bit worse
+        since there are up to N * J extra samples with IterableDatasets.
+      * The sharding (splitting of dataset into TFRecord) files imposes limitations on the number of
+        replicas and dataloader workers you can use. For really small datasets that only contain a few shards
+        you may have to train non-distributed w/ 1-2 dataloader workers. This is likely not a huge concern as the
+        benefit of distributed training or fast dataloading should be much less for small datasets.
+      * This wrapper is currently configured to return individual, decompressed image samples from the TFDS
+        dataset. The augmentation (transforms) and batching is still done in PyTorch. It would be possible
+        to specify TF augmentation fn and return augmented batches w/ some modifications to other downstream
+        components.
+
+    """
+
+    def __init__(
+            self,
+            name,
+            root=None,
+            split='train',
+            class_map=None,
+            is_training=False,
+            batch_size=1,
+            download=False,
+            repeats=0,
+            seed=42,
+            input_key='image',
+            input_img_mode='RGB',
+            target_key='label',
+            target_img_mode='',
+            prefetch_size=None,
+            shuffle_size=None,
+            max_threadpool_size=None
+    ):
+        """ Tensorflow-datasets Wrapper
+
+        Args:
+            root: root data dir (ie your TFDS_DATA_DIR. not dataset specific sub-dir)
+            name: tfds dataset name (eg `imagenet2012`)
+            split: tfds dataset split (can use all TFDS split strings eg `train[:10%]`)
+            is_training: training mode, shuffle enabled, dataset len rounded by batch_size
+            batch_size: batch_size to use to unsure total samples % batch_size == 0 in training across all dis nodes
+            download: download and build TFDS dataset if set, otherwise must use tfds CLI
+            repeats: iterate through (repeat) the dataset this many times per iteration (once if 0 or 1)
+            seed: common seed for shard shuffle across all distributed/worker instances
+            input_key: name of Feature to return as data (input)
+            input_img_mode: image mode if input is an image (currently PIL mode string)
+            target_key: name of Feature to return as target (label)
+            target_img_mode: image mode if target is an image (currently PIL mode string)
+            prefetch_size: override default tf.data prefetch buffer size
+            shuffle_size: override default tf.data shuffle buffer size
+            max_threadpool_size: override default threadpool size for tf.data
+        """
+        super().__init__()
+        self.root = root
+        self.split = split
+        self.is_training = is_training
+        self.batch_size = batch_size
+        self.repeats = repeats
+        self.common_seed = seed  # a seed that's fixed across all worker / distributed instances
+
+        # performance settings
+        self.prefetch_size = prefetch_size or PREFETCH_SIZE
+        self.shuffle_size = shuffle_size or SHUFFLE_SIZE
+        self.max_threadpool_size = max_threadpool_size or MAX_TP_SIZE
+
+        # TFDS builder and split information
+        self.input_key = input_key  # FIXME support tuples / lists of inputs and targets and full range of Feature
+        self.input_img_mode = input_img_mode
+        self.target_key = target_key
+        self.target_img_mode = target_img_mode  # for dense pixel targets
+        self.builder = tfds.builder(name, data_dir=root)
+        # NOTE: the tfds command line app can be used download & prepare datasets if you don't enable download flag
+        if download:
+            self.builder.download_and_prepare()
+        self.remap_class = False
+        if class_map:
+            self.class_to_idx = load_class_map(class_map)
+            self.remap_class = True
+        else:
+            self.class_to_idx = get_class_labels(self.builder.info) if self.target_key == 'label' else {}
+        self.split_info = self.builder.info.splits[split]
+        self.num_samples = self.split_info.num_examples
+
+        # Distributed world state
+        self.dist_rank = 0
+        self.dist_num_replicas = 1
+        if dist.is_available() and dist.is_initialized() and dist.get_world_size() > 1:
+            self.dist_rank = dist.get_rank()
+            self.dist_num_replicas = dist.get_world_size()
+
+        # Attributes that are updated in _lazy_init, including the tf.data pipeline itself
+        self.global_num_workers = 1
+        self.num_workers = 1
+        self.worker_info = None
+        self.worker_seed = 0  # seed unique to each work instance
+        self.subsplit = None  # set when data is distributed across workers using sub-splits
+        self.ds = None  # initialized lazily on each dataloader worker process
+        self.init_count = 0  # number of ds TF data pipeline initializations
+        self.epoch_count = SharedCount()
+        # FIXME need to determine if reinit_each_iter is necessary. I'm don't completely trust behaviour
+        #  of `shuffle_reshuffle_each_iteration` when there are multiple workers / nodes across epochs
+        self.reinit_each_iter = self.is_training
+
+    def set_epoch(self, count):
+        self.epoch_count.value = count
+
+    def set_loader_cfg(
+            self,
+            num_workers: Optional[int] = None,
+    ):
+        if self.ds is not None:
+            return
+        if num_workers is not None:
+            self.num_workers = num_workers
+            self.global_num_workers = self.dist_num_replicas * self.num_workers
+
+    def _lazy_init(self):
+        """ Lazily initialize the dataset.
+
+        This is necessary to init the Tensorflow dataset pipeline in the (dataloader) process that
+        will be using the dataset instance. The __init__ method is called on the main process,
+        this will be called in a dataloader worker process.
+
+        NOTE: There will be problems if you try to re-use this dataset across different loader/worker
+        instances once it has been initialized. Do not call any dataset methods that can call _lazy_init
+        before it is passed to dataloader.
+        """
+        worker_info = torch.utils.data.get_worker_info()
+
+        # setup input context to split dataset across distributed processes
+        num_workers = 1
+        global_worker_id = 0
+        if worker_info is not None:
+            self.worker_info = worker_info
+            self.worker_seed = worker_info.seed
+            self.num_workers = worker_info.num_workers
+            self.global_num_workers = self.dist_num_replicas * self.num_workers
+            global_worker_id = self.dist_rank * self.num_workers + worker_info.id
+
+            """ Data sharding
+            InputContext will assign subset of underlying TFRecord files to each 'pipeline' if used.
+            My understanding is that using split, the underling TFRecord files will shuffle (shuffle_files=True)
+            between the splits each iteration, but that understanding could be wrong.
+
+            I am currently using a mix of InputContext shard assignment and fine-grained sub-splits for distributing
+            the data across workers. For training InputContext is used to assign shards to nodes unless num_shards
+            in dataset < total number of workers. Otherwise sub-split API is used for datasets without enough shards or
+            for validation where we can't drop samples and need to avoid minimize uneven splits to avoid padding.
+            """
+            should_subsplit = self.global_num_workers > 1 and (
+                    self.split_info.num_shards < self.global_num_workers or not self.is_training)
+            if should_subsplit:
+                # split the dataset w/o using sharding for more even samples / worker, can result in less optimal
+                # read patterns for distributed training (overlap across shards) so better to use InputContext there
+                if has_buggy_even_splits:
+                    # my even_split workaround doesn't work on subsplits, upgrade tfds!
+                    if not isinstance(self.split_info, tfds.core.splits.SubSplitInfo):
+                        subsplits = even_split_indices(self.split, self.global_num_workers, self.num_samples)
+                        self.subsplit = subsplits[global_worker_id]
+                else:
+                    subsplits = tfds.even_splits(self.split, self.global_num_workers)
+                    self.subsplit = subsplits[global_worker_id]
+
+        input_context = None
+        if self.global_num_workers > 1 and self.subsplit is None:
+            # set input context to divide shards among distributed replicas
+            input_context = tf.distribute.InputContext(
+                num_input_pipelines=self.global_num_workers,
+                input_pipeline_id=global_worker_id,
+                num_replicas_in_sync=self.dist_num_replicas  # FIXME does this arg have any impact?
+            )
+        read_config = tfds.ReadConfig(
+            shuffle_seed=self.common_seed + self.epoch_count.value,
+            shuffle_reshuffle_each_iteration=True,
+            input_context=input_context,
+        )
+        ds = self.builder.as_dataset(
+            split=self.subsplit or self.split,
+            shuffle_files=self.is_training,
+            decoders=dict(image=decode_example(channels=1 if self.input_img_mode == 'L' else 3)),
+            read_config=read_config,
+        )
+        # avoid overloading threading w/ combo of TF ds threads + PyTorch workers
+        options = tf.data.Options()
+        thread_member = 'threading' if hasattr(options, 'threading') else 'experimental_threading'
+        getattr(options, thread_member).private_threadpool_size = max(1, self.max_threadpool_size // self.num_workers)
+        getattr(options, thread_member).max_intra_op_parallelism = 1
+        ds = ds.with_options(options)
+        if self.is_training or self.repeats > 1:
+            # to prevent excessive drop_last batch behaviour w/ IterableDatasets
+            # see warnings at https://pytorch.org/docs/stable/data.html#multi-process-data-loading
+            ds = ds.repeat()  # allow wrap around and break iteration manually
+        if self.is_training:
+            ds = ds.shuffle(min(self.num_samples, self.shuffle_size) // self.global_num_workers, seed=self.worker_seed)
+        ds = ds.prefetch(min(self.num_samples // self.global_num_workers, self.prefetch_size))
+        self.ds = tfds.as_numpy(ds)
+        self.init_count += 1
+
+    def _num_samples_per_worker(self):
+        num_worker_samples = \
+            max(1, self.repeats) * self.num_samples / max(self.global_num_workers, self.dist_num_replicas)
+        if self.is_training or self.dist_num_replicas > 1:
+            num_worker_samples = math.ceil(num_worker_samples)
+        if self.is_training:
+            num_worker_samples = math.ceil(num_worker_samples / self.batch_size) * self.batch_size
+        return int(num_worker_samples)
+
+    def __iter__(self):
+        if self.ds is None or self.reinit_each_iter:
+            self._lazy_init()
+
+        # Compute a rounded up sample count that is used to:
+        #   1. make batches even cross workers & replicas in distributed validation.
+        #     This adds extra samples and will slightly alter validation results.
+        #   2. determine loop ending condition in training w/ repeat enabled so that only full batch_size
+        #     batches are produced (underlying tfds iter wraps around)
+        target_sample_count = self._num_samples_per_worker()
+
+        # Iterate until exhausted or sample count hits target when training (ds.repeat enabled)
+        sample_count = 0
+        for sample in self.ds:
+            input_data = sample[self.input_key]
+            if self.input_img_mode:
+                if self.input_img_mode == 'L' and input_data.ndim == 3:
+                    input_data = input_data[:, :, 0]
+                input_data = Image.fromarray(input_data, mode=self.input_img_mode)
+            target_data = sample[self.target_key]
+            if self.target_img_mode:
+                # dense pixel target
+                target_data = Image.fromarray(target_data, mode=self.target_img_mode)
+            elif self.remap_class:
+                target_data = self.class_to_idx[target_data]
+            yield input_data, target_data
+            sample_count += 1
+            if self.is_training and sample_count >= target_sample_count:
+                # Need to break out of loop when repeat() is enabled for training w/ oversampling
+                # this results in extra samples per epoch but seems more desirable than dropping
+                # up to N*J batches per epoch (where N = num distributed processes, and J = num worker processes)
+                break
+
+        # Pad across distributed nodes (make counts equal by adding samples)
+        if not self.is_training and self.dist_num_replicas > 1 and self.subsplit is not None and \
+                0 < sample_count < target_sample_count:
+            # Validation batch padding only done for distributed training where results are reduced across nodes.
+            # For single process case, it won't matter if workers return different batch sizes.
+            # If using input_context or % based splits, sample count can vary significantly across workers and this
+            # approach should not be used (hence disabled if self.subsplit isn't set).
+            while sample_count < target_sample_count:
+                yield input_data, target_data  # yield prev sample again
+                sample_count += 1
+
+    def __len__(self):
+        num_samples = self._num_samples_per_worker() * self.num_workers
+        return num_samples
+
+    def _filename(self, index, basename=False, absolute=False):
+        assert False, "Not supported"  # no random access to samples
+
+    def filenames(self, basename=False, absolute=False):
+        """ Return all filenames in dataset, overrides base"""
+        if self.ds is None:
+            self._lazy_init()
+        names = []
+        for sample in self.ds:
+            if len(names) > self.num_samples:
+                break  # safety for ds.repeat() case
+            if 'file_name' in sample:
+                name = sample['file_name']
+            elif 'filename' in sample:
+                name = sample['filename']
+            elif 'id' in sample:
+                name = sample['id']
+            else:
+                assert False, "No supported name field present"
+            names.append(name)
+        return names

timm/data/readers/reader_wds.py

@@ -0,0 +1,466 @@
+""" Dataset reader for webdataset
+
+Hacked together by / Copyright 2022 Ross Wightman
+"""
+import io
+import json
+import logging
+import math
+import os
+import random
+import sys
+from dataclasses import dataclass
+from functools import partial
+from itertools import islice
+from typing import Any, Callable, Dict, List, Optional, Tuple
+
+import torch
+import torch.distributed as dist
+import yaml
+from PIL import Image
+from torch.utils.data import Dataset, IterableDataset, get_worker_info
+
+try:
+    import webdataset as wds
+    from webdataset.filters import _shuffle, getfirst
+    from webdataset.shardlists import expand_urls
+    from webdataset.tariterators import base_plus_ext, url_opener, tar_file_expander, valid_sample
+except ImportError:
+    wds = None
+    expand_urls = None
+
+from .class_map import load_class_map
+from .reader import Reader
+from .shared_count import SharedCount
+
+_logger = logging.getLogger(__name__)
+
+SAMPLE_SHUFFLE_SIZE = int(os.environ.get('WDS_SHUFFLE_SIZE', 8192))
+SAMPLE_INITIAL_SIZE = int(os.environ.get('WDS_INITIAL_SIZE', 2048))
+
+
+def _load_info(root, names=('_info.json', 'info.json')):
+    if isinstance(names, str):
+        names = (names,)
+    tried = []
+    err_str = ''
+    for n in names:
+        full_path = os.path.join(root, n)
+        try:
+            tried.append(full_path)
+            with wds.gopen(full_path) as f:
+                if n.endswith('.json'):
+                    info_dict = json.load(f)
+                else:
+                    info_dict = yaml.safe_load(f)
+            return info_dict
+        except Exception as e:
+            err_str = str(e)
+
+    _logger.warning(
+        f'Dataset info file not found at {tried}. Error: {err_str}. '
+        'Falling back to provided split and size arg.')
+    return {}
+
+
+@dataclass
+class SplitInfo:
+    num_samples: int
+    filenames: Tuple[str]
+    shard_lengths: Tuple[int] = ()
+    alt_label: str = ''
+    name: str = ''
+
+
+def _parse_split_info(split: str, info: Dict):
+    def _info_convert(dict_info):
+        return SplitInfo(
+            num_samples=dict_info['num_samples'],
+            filenames=tuple(dict_info['filenames']),
+            shard_lengths=tuple(dict_info['shard_lengths']),
+            alt_label=dict_info.get('alt_label', ''),
+            name=dict_info['name'],
+        )
+
+    if 'tar' in split or '..' in split:
+        # split in WDS string braceexpand format, sample count can be included with a | separator
+        # ex: `dataset-split-{0000..9999}.tar|100000` for 9999 shards, covering 100,000 samples
+        split = split.split('|')
+        num_samples = 0
+        split_name = ''
+        if len(split) > 1:
+            num_samples = int(split[1])
+        split = split[0]
+        if '::' not in split:
+            split_parts = split.split('-', 3)
+            split_idx = len(split_parts) - 1
+            if split_idx and 'splits' in info and split_parts[split_idx] in info['splits']:
+                split_name = split_parts[split_idx]
+
+        split_filenames = expand_urls(split)
+        if split_name:
+            split_info = info['splits'][split_name]
+            if not num_samples:
+                _fc = {f: c for f, c in zip(split_info['filenames'], split_info['shard_lengths'])}
+                num_samples = sum(_fc[f] for f in split_filenames)
+                split_info['filenames'] = tuple(_fc.keys())
+                split_info['shard_lengths'] = tuple(_fc.values())
+                split_info['num_samples'] = num_samples
+            split_info = _info_convert(split_info)
+        else:
+            split_info = SplitInfo(
+                name=split_name,
+                num_samples=num_samples,
+                filenames=split_filenames,
+            )
+    else:
+        if 'splits' not in info or split not in info['splits']:
+            raise RuntimeError(f"split {split} not found in info ({info.get('splits', {}).keys()})")
+        split = split
+        split_info = info['splits'][split]
+        split_info = _info_convert(split_info)
+
+    return split_info
+
+
+def log_and_continue(exn):
+    """Call in an exception handler to ignore exceptions, isssue a warning, and continue."""
+    _logger.warning(f'Handling webdataset error ({repr(exn)}). Ignoring.')
+    # NOTE: try force an exit on errors that are clearly code / config and not transient
+    if isinstance(exn, TypeError):
+        raise exn
+    return True
+
+
+def _decode(
+        sample,
+        image_key='jpg',
+        image_mode='RGB',
+        target_key='cls',
+        alt_label=''
+):
+    """ Custom sample decode
+    * decode and convert PIL Image
+    * cls byte string label to int
+    * pass through JSON byte string (if it exists) without parse
+    """
+    # decode class label, skip if alternate label not valid
+    if alt_label:
+        # alternative labels are encoded in json metadata
+        meta = json.loads(sample['json'])
+        class_label = int(meta[alt_label])
+        if class_label < 0:
+            # skipped labels currently encoded as -1, may change to a null/None value
+            return None
+    else:
+        class_label = int(sample[target_key])
+
+    # decode image
+    img = getfirst(sample, image_key)
+    with io.BytesIO(img) as b:
+        img = Image.open(b)
+        img.load()
+    if image_mode:
+        img = img.convert(image_mode)
+
+    # json passed through in undecoded state
+    decoded = dict(jpg=img, cls=class_label, json=sample.get('json', None))
+    return decoded
+
+
+def pytorch_worker_seed():
+    """get dataloader worker seed from pytorch"""
+    worker_info = get_worker_info()
+    if worker_info is not None:
+        # favour the seed already created for pytorch dataloader workers if it exists
+        return worker_info.seed
+    # fallback to wds rank based seed
+    return wds.utils.pytorch_worker_seed()
+
+
+if wds is not None:
+    # conditional to avoid mandatory wds import (via inheritance of wds.PipelineStage)
+
+    class detshuffle2(wds.PipelineStage):
+        def __init__(
+                self,
+                bufsize=1000,
+                initial=100,
+                seed=0,
+                epoch=-1,
+        ):
+            self.bufsize = bufsize
+            self.initial = initial
+            self.seed = seed
+            self.epoch = epoch
+
+        def run(self, src):
+            if isinstance(self.epoch, SharedCount):
+                epoch = self.epoch.value
+            else:
+                # NOTE: this is epoch tracking is problematic in a multiprocess (dataloader workers or train)
+                # situation as different workers may wrap at different times (or not at all).
+                self.epoch += 1
+                epoch = self.epoch
+
+            if self.seed < 0:
+                seed = pytorch_worker_seed() + epoch
+            else:
+                seed = self.seed + epoch
+            # _logger.info(f'shuffle seed: {self.seed}, {seed}, epoch: {epoch}')  # FIXME temporary
+            rng = random.Random(seed)
+            return _shuffle(src, self.bufsize, self.initial, rng)
+
+else:
+    detshuffle2 = None
+
+
+class ResampledShards2(IterableDataset):
+    """An iterable dataset yielding a list of urls."""
+
+    def __init__(
+        self,
+        urls,
+        nshards=sys.maxsize,
+        worker_seed=None,
+        deterministic=True,
+        epoch=-1,
+    ):
+        """Sample shards from the shard list with replacement.
+
+        :param urls: a list of URLs as a Python list or brace notation string
+        """
+        super().__init__()
+        urls = wds.shardlists.expand_urls(urls)
+        self.urls = urls
+        assert isinstance(self.urls[0], str)
+        self.nshards = nshards
+        self.rng = random.Random()
+        self.worker_seed = pytorch_worker_seed if worker_seed is None else worker_seed
+        self.deterministic = deterministic
+        self.epoch = epoch
+
+    def __iter__(self):
+        """Return an iterator over the shards."""
+        if isinstance(self.epoch, SharedCount):
+            epoch = self.epoch.value
+        else:
+            # NOTE: this is epoch tracking is problematic in a multiprocess (dataloader workers or train)
+            # situation as different workers may wrap at different times (or not at all).
+            self.epoch += 1
+            epoch = self.epoch
+
+        if self.deterministic:
+            # reset seed w/ epoch if deterministic, worker seed should be deterministic due to arg.seed
+            self.rng = random.Random(self.worker_seed() + epoch)
+
+        for _ in range(self.nshards):
+            index = self.rng.randint(0, len(self.urls) - 1)
+            yield dict(url=self.urls[index])
+
+
+class ReaderWds(Reader):
+    def __init__(
+            self,
+            root: str,
+            name: Optional[str] = None,
+            split: str = 'train',
+            is_training: bool = False,
+            num_samples: Optional[int] = None,
+            batch_size: int = 1,
+            repeats: int = 0,
+            seed: int = 42,
+            class_map: Optional[dict] = None,
+            input_key: str = 'jpg;png;webp',
+            input_img_mode: str = 'RGB',
+            target_key: str = 'cls',
+            target_img_mode: str = '',
+            filename_key: str = 'filename',
+            sample_shuffle_size: Optional[int] = None,
+            smaple_initial_size: Optional[int] = None,
+    ):
+        super().__init__()
+        if wds is None:
+            raise RuntimeError(
+                'Please install webdataset 0.2.x package `pip install git+https://github.com/webdataset/webdataset`.')
+        self.root = root
+        self.is_training = is_training
+        self.batch_size = batch_size
+        self.repeats = repeats
+        self.common_seed = seed  # a seed that's fixed across all worker / distributed instances
+        self.shard_shuffle_size = 500
+        self.sample_shuffle_size = sample_shuffle_size or SAMPLE_SHUFFLE_SIZE
+        self.sample_initial_size = smaple_initial_size or SAMPLE_INITIAL_SIZE
+
+        self.input_key = input_key
+        self.input_img_mode = input_img_mode
+        self.target_key = target_key
+        self.filename_key = filename_key
+        self.key_ext = '.JPEG'  # extension to add to key for original filenames (DS specific, default ImageNet)
+
+        self.info = _load_info(self.root)
+        self.split_info = _parse_split_info(split, self.info)
+        if num_samples is not None:
+            self.num_samples = num_samples
+        else:
+            self.num_samples = self.split_info.num_samples
+        if not self.num_samples:
+            raise RuntimeError(f'Invalid split definition, num_samples not specified.')
+        self.remap_class = False
+        if class_map:
+            self.class_to_idx = load_class_map(class_map)
+            self.remap_class = True
+        else:
+            self.class_to_idx = {}
+
+        # Distributed world state
+        self.dist_rank = 0
+        self.dist_num_replicas = 1
+        if dist.is_available() and dist.is_initialized() and dist.get_world_size() > 1:
+            self.dist_rank = dist.get_rank()
+            self.dist_num_replicas = dist.get_world_size()
+
+        # Attributes that are updated in _lazy_init
+        self.worker_info = None
+        self.worker_id = 0
+        self.worker_seed = seed  # seed unique to each worker instance
+        self.num_workers = 1
+        self.global_worker_id = 0
+        self.global_num_workers = 1
+        self.init_count = 0
+        self.epoch_count = SharedCount()
+
+        # DataPipeline is lazy init, the majority of WDS DataPipeline could be init here, BUT, shuffle seed
+        # is not handled in manner where it can be deterministic for each worker AND initialized up front
+        self.ds = None
+
+    def set_epoch(self, count):
+        self.epoch_count.value = count
+
+    def set_loader_cfg(
+            self,
+            num_workers: Optional[int] = None,
+    ):
+        if self.ds is not None:
+            return
+        if num_workers is not None:
+            self.num_workers = num_workers
+            self.global_num_workers = self.dist_num_replicas * self.num_workers
+
+    def _lazy_init(self):
+        """ Lazily initialize worker (in worker processes)
+        """
+        if self.worker_info is None:
+            worker_info = torch.utils.data.get_worker_info()
+            if worker_info is not None:
+                self.worker_info = worker_info
+                self.worker_id = worker_info.id
+                self.worker_seed = worker_info.seed
+                self.num_workers = worker_info.num_workers
+            self.global_num_workers = self.dist_num_replicas * self.num_workers
+            self.global_worker_id = self.dist_rank * self.num_workers + self.worker_id
+
+        # init data pipeline
+        abs_shard_filenames = [os.path.join(self.root, f) for f in self.split_info.filenames]
+        pipeline = [wds.SimpleShardList(abs_shard_filenames)]
+        # at this point we have an iterator over all the shards
+        if self.is_training:
+            pipeline.extend([
+                detshuffle2(
+                    self.shard_shuffle_size,
+                    seed=self.common_seed,
+                    epoch=self.epoch_count,
+                ),
+                self._split_by_node_and_worker,
+                # at this point, we have an iterator over the shards assigned to each worker
+                wds.tarfile_to_samples(handler=log_and_continue),
+                wds.shuffle(
+                    bufsize=self.sample_shuffle_size,
+                    initial=self.sample_initial_size,
+                    rng=random.Random(self.worker_seed) # this is why we lazy-init whole DataPipeline
+                ),
+            ])
+        else:
+            pipeline.extend([
+                self._split_by_node_and_worker,
+                # at this point, we have an iterator over the shards assigned to each worker
+                wds.tarfile_to_samples(handler=log_and_continue),
+            ])
+        pipeline.extend([
+            wds.map(
+                partial(
+                    _decode,
+                    image_key=self.input_key,
+                    image_mode=self.input_img_mode,
+                    alt_label=self.split_info.alt_label,
+                ),
+                handler=log_and_continue,
+            ),
+            wds.rename(image=self.input_key, target=self.target_key)
+        ])
+        self.ds = wds.DataPipeline(*pipeline)
+
+    def _split_by_node_and_worker(self, src):
+        if self.global_num_workers > 1:
+            for s in islice(src, self.global_worker_id, None, self.global_num_workers):
+                yield s
+        else:
+            for s in src:
+                yield s
+
+    def _num_samples_per_worker(self):
+        num_worker_samples = self.num_samples / max(self.global_num_workers, self.dist_num_replicas)
+        if self.is_training or self.dist_num_replicas > 1:
+            num_worker_samples = math.ceil(num_worker_samples)
+        if self.is_training:
+            num_worker_samples = math.ceil(num_worker_samples / self.batch_size) * self.batch_size
+        return int(num_worker_samples)
+
+    def __iter__(self):
+        if self.ds is None:
+            self._lazy_init()
+
+        num_worker_samples = self._num_samples_per_worker()
+        if self.is_training or self.dist_num_replicas > 1:
+            # NOTE: doing distributed validation w/ WDS is messy, hard to meet constraints that
+            # same # of batches needed across all replicas w/ seeing each sample once.
+            # with_epoch() is simple but could miss a shard's worth of samples in some workers,
+            # and duplicate in others. Best to keep num DL workers low and a divisor of #val shards.
+            ds = self.ds.with_epoch(num_worker_samples)
+        else:
+            ds = self.ds
+
+        i = 0
+        # _logger.info(f'start {i}, {self.worker_id}')  # FIXME temporary debug
+        for sample in ds:
+            target = sample['target']
+            if self.remap_class:
+                target = self.class_to_idx[target]
+            yield sample['image'], target
+            i += 1
+        # _logger.info(f'end {i}, {self.worker_id}')  # FIXME temporary debug
+
+    def __len__(self):
+        num_samples = self._num_samples_per_worker() * self.num_workers
+        return num_samples
+
+    def _filename(self, index, basename=False, absolute=False):
+        assert False, "Not supported"  # no random access to examples
+
+    def filenames(self, basename=False, absolute=False):
+        """ Return all filenames in dataset, overrides base"""
+        if self.ds is None:
+            self._lazy_init()
+
+        names = []
+        for sample in self.ds:
+            if self.filename_key in sample:
+                name = sample[self.filename_key]
+            elif '__key__' in sample:
+                name = sample['__key__'] + self.key_ext
+            else:
+                assert False, "No supported name field present"
+            names.append(name)
+            if len(names) >= self.num_samples:
+                break  # safety for ds.repeat() case
+        return names

timm/data/readers/shared_count.py

@@ -0,0 +1,14 @@
+from multiprocessing import Value
+
+
+class SharedCount:
+    def __init__(self, epoch: int = 0):
+        self.shared_epoch = Value('i', epoch)
+
+    @property
+    def value(self):
+        return self.shared_epoch.value
+
+    @value.setter
+    def value(self, epoch):
+        self.shared_epoch.value = epoch

timm/data/real_labels.py

@@ -0,0 +1,47 @@
+""" Real labels evaluator for ImageNet
+Paper: `Are we done with ImageNet?` - https://arxiv.org/abs/2006.07159
+Based on Numpy example at https://github.com/google-research/reassessed-imagenet
+
+Hacked together by / Copyright 2020 Ross Wightman
+"""
+import os
+import json
+import numpy as np
+import pkgutil
+
+
+class RealLabelsImagenet:
+
+    def __init__(self, filenames, real_json=None, topk=(1, 5)):
+        if real_json is not None:
+            with open(real_json) as real_labels:
+                real_labels = json.load(real_labels)
+        else:
+            real_labels = json.loads(
+                pkgutil.get_data(__name__, os.path.join('_info', 'imagenet_real_labels.json')).decode('utf-8'))
+        real_labels = {f'ILSVRC2012_val_{i + 1:08d}.JPEG': labels for i, labels in enumerate(real_labels)}
+        self.real_labels = real_labels
+        self.filenames = filenames
+        assert len(self.filenames) == len(self.real_labels)
+        self.topk = topk
+        self.is_correct = {k: [] for k in topk}
+        self.sample_idx = 0
+
+    def add_result(self, output):
+        maxk = max(self.topk)
+        _, pred_batch = output.topk(maxk, 1, True, True)
+        pred_batch = pred_batch.cpu().numpy()
+        for pred in pred_batch:
+            filename = self.filenames[self.sample_idx]
+            filename = os.path.basename(filename)
+            if self.real_labels[filename]:
+                for k in self.topk:
+                    self.is_correct[k].append(
+                        any([p in self.real_labels[filename] for p in pred[:k]]))
+            self.sample_idx += 1
+
+    def get_accuracy(self, k=None):
+        if k is None:
+            return {k: float(np.mean(self.is_correct[k])) * 100 for k in self.topk}
+        else:
+            return float(np.mean(self.is_correct[k])) * 100

timm/data/tf_preprocessing.py

@@ -0,0 +1,233 @@
+""" Tensorflow Preprocessing Adapter
+
+Allows use of Tensorflow preprocessing pipeline in PyTorch Transform
+
+Copyright of original Tensorflow code below.
+
+Hacked together by / Copyright 2020 Ross Wightman
+"""
+
+# Copyright 2018 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""ImageNet preprocessing for MnasNet."""
+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,
+                                min_object_covered=0.1,
+                                aspect_ratio_range=(0.75, 1.33),
+                                area_range=(0.05, 1.0),
+                                max_attempts=100,
+                                scope=None):
+    """Generates cropped_image using one of the bboxes randomly distorted.
+
+    See `tf.image.sample_distorted_bounding_box` for more documentation.
+
+    Args:
+      image_bytes: `Tensor` of binary image data.
+      bbox: `Tensor` of bounding boxes arranged `[1, num_boxes, coords]`
+          where each coordinate is [0, 1) and the coordinates are arranged
+          as `[ymin, xmin, ymax, xmax]`. If num_boxes is 0 then use the whole
+          image.
+      min_object_covered: An optional `float`. Defaults to `0.1`. The cropped
+          area of the image must contain at least this fraction of any bounding
+          box supplied.
+      aspect_ratio_range: An optional list of `float`s. The cropped area of the
+          image must have an aspect ratio = width / height within this range.
+      area_range: An optional list of `float`s. The cropped area of the image
+          must contain a fraction of the supplied image within in this range.
+      max_attempts: An optional `int`. Number of attempts at generating a cropped
+          region of the image of the specified constraints. After `max_attempts`
+          failures, return the entire image.
+      scope: Optional `str` for name scope.
+    Returns:
+      cropped image `Tensor`
+    """
+    with tf.name_scope(scope, 'distorted_bounding_box_crop', [image_bytes, bbox]):
+        shape = tf.image.extract_jpeg_shape(image_bytes)
+        sample_distorted_bounding_box = tf.image.sample_distorted_bounding_box(
+            shape,
+            bounding_boxes=bbox,
+            min_object_covered=min_object_covered,
+            aspect_ratio_range=aspect_ratio_range,
+            area_range=area_range,
+            max_attempts=max_attempts,
+            use_image_if_no_bounding_boxes=True)
+        bbox_begin, bbox_size, _ = sample_distorted_bounding_box
+
+        # Crop the image to the specified bounding box.
+        offset_y, offset_x, _ = tf.unstack(bbox_begin)
+        target_height, target_width, _ = tf.unstack(bbox_size)
+        crop_window = tf.stack([offset_y, offset_x, target_height, target_width])
+        image = tf.image.decode_and_crop_jpeg(image_bytes, crop_window, channels=3)
+
+        return image
+
+
+def _at_least_x_are_equal(a, b, x):
+    """At least `x` of `a` and `b` `Tensors` are equal."""
+    match = tf.equal(a, b)
+    match = tf.cast(match, tf.int32)
+    return tf.greater_equal(tf.reduce_sum(match), x)
+
+
+def _decode_and_random_crop(image_bytes, image_size, resize_method):
+    """Make a random crop of image_size."""
+    bbox = tf.constant([0.0, 0.0, 1.0, 1.0], dtype=tf.float32, shape=[1, 1, 4])
+    image = distorted_bounding_box_crop(
+        image_bytes,
+        bbox,
+        min_object_covered=0.1,
+        aspect_ratio_range=(3. / 4, 4. / 3.),
+        area_range=(0.08, 1.0),
+        max_attempts=10,
+        scope=None)
+    original_shape = tf.image.extract_jpeg_shape(image_bytes)
+    bad = _at_least_x_are_equal(original_shape, tf.shape(image), 3)
+
+    image = tf.cond(
+        bad,
+        lambda: _decode_and_center_crop(image_bytes, image_size),
+        lambda: tf.image.resize([image], [image_size, image_size], resize_method)[0])
+
+    return image
+
+
+def _decode_and_center_crop(image_bytes, image_size, resize_method):
+    """Crops to center of image with padding then scales image_size."""
+    shape = tf.image.extract_jpeg_shape(image_bytes)
+    image_height = shape[0]
+    image_width = shape[1]
+
+    padded_center_crop_size = tf.cast(
+        ((image_size / (image_size + CROP_PADDING)) *
+         tf.cast(tf.minimum(image_height, image_width), tf.float32)),
+        tf.int32)
+
+    offset_height = ((image_height - padded_center_crop_size) + 1) // 2
+    offset_width = ((image_width - padded_center_crop_size) + 1) // 2
+    crop_window = tf.stack([offset_height, offset_width,
+                            padded_center_crop_size, padded_center_crop_size])
+    image = tf.image.decode_and_crop_jpeg(image_bytes, crop_window, channels=3)
+    image = tf.image.resize([image], [image_size, image_size], resize_method)[0]
+
+    return image
+
+
+def _flip(image):
+    """Random horizontal image flip."""
+    image = tf.image.random_flip_left_right(image)
+    return image
+
+
+def preprocess_for_train(image_bytes, use_bfloat16, image_size=IMAGE_SIZE, interpolation='bicubic'):
+    """Preprocesses the given image for evaluation.
+
+    Args:
+      image_bytes: `Tensor` representing an image binary of arbitrary size.
+      use_bfloat16: `bool` for whether to use bfloat16.
+      image_size: image size.
+      interpolation: image interpolation method
+
+    Returns:
+      A preprocessed image `Tensor`.
+    """
+    resize_method = tf.image.ResizeMethod.BICUBIC if interpolation == 'bicubic' else tf.image.ResizeMethod.BILINEAR
+    image = _decode_and_random_crop(image_bytes, image_size, resize_method)
+    image = _flip(image)
+    image = tf.reshape(image, [image_size, image_size, 3])
+    image = tf.image.convert_image_dtype(
+        image, dtype=tf.bfloat16 if use_bfloat16 else tf.float32)
+    return image
+
+
+def preprocess_for_eval(image_bytes, use_bfloat16, image_size=IMAGE_SIZE, interpolation='bicubic'):
+    """Preprocesses the given image for evaluation.
+
+    Args:
+      image_bytes: `Tensor` representing an image binary of arbitrary size.
+      use_bfloat16: `bool` for whether to use bfloat16.
+      image_size: image size.
+      interpolation: image interpolation method
+
+    Returns:
+      A preprocessed image `Tensor`.
+    """
+    resize_method = tf.image.ResizeMethod.BICUBIC if interpolation == 'bicubic' else tf.image.ResizeMethod.BILINEAR
+    image = _decode_and_center_crop(image_bytes, image_size, resize_method)
+    image = tf.reshape(image, [image_size, image_size, 3])
+    image = tf.image.convert_image_dtype(
+        image, dtype=tf.bfloat16 if use_bfloat16 else tf.float32)
+    return image
+
+
+def preprocess_image(image_bytes,
+                     is_training=False,
+                     use_bfloat16=False,
+                     image_size=IMAGE_SIZE,
+                     interpolation='bicubic'):
+    """Preprocesses the given image.
+
+    Args:
+      image_bytes: `Tensor` representing an image binary of arbitrary size.
+      is_training: `bool` for whether the preprocessing is for training.
+      use_bfloat16: `bool` for whether to use bfloat16.
+      image_size: image size.
+      interpolation: image interpolation method
+
+    Returns:
+      A preprocessed image `Tensor` with value range of [0, 255].
+    """
+    if is_training:
+        return preprocess_for_train(image_bytes, use_bfloat16, image_size, interpolation)
+    else:
+        return preprocess_for_eval(image_bytes, use_bfloat16, image_size, interpolation)
+
+
+class TfPreprocessTransform:
+
+    def __init__(self, is_training=False, size=224, interpolation='bicubic'):
+        self.is_training = is_training
+        self.size = size[0] if isinstance(size, tuple) else size
+        self.interpolation = interpolation
+        self._image_bytes = None
+        self.process_image = self._build_tf_graph()
+        self.sess = None
+
+    def _build_tf_graph(self):
+        with tf.device('/cpu:0'):
+            self._image_bytes = tf.placeholder(
+                shape=[],
+                dtype=tf.string,
+            )
+            img = preprocess_image(
+                self._image_bytes, self.is_training, False, self.size, self.interpolation)
+        return img
+
+    def __call__(self, image_bytes):
+        if self.sess is None:
+            self.sess = tf.Session()
+        img = self.sess.run(self.process_image, feed_dict={self._image_bytes: image_bytes})
+        img = img.round().clip(0, 255).astype(np.uint8)
+        if img.ndim < 3:
+            img = np.expand_dims(img, axis=-1)
+        img = np.rollaxis(img, 2)  # HWC to CHW
+        return img

timm/data/transforms.py

@@ -0,0 +1,534 @@
+import math
+import numbers
+import random
+import warnings
+from typing import List, Sequence, Tuple, Union
+
+import torch
+import torchvision.transforms.functional as F
+try:
+    from torchvision.transforms.functional import InterpolationMode
+    has_interpolation_mode = True
+except ImportError:
+    has_interpolation_mode = False
+from PIL import Image
+import numpy as np
+
+__all__ = [
+    "ToNumpy", "ToTensor", "str_to_interp_mode", "str_to_pil_interp", "interp_mode_to_str",
+    "RandomResizedCropAndInterpolation", "CenterCropOrPad", "center_crop_or_pad", "crop_or_pad",
+    "RandomCropOrPad", "RandomPad", "ResizeKeepRatio", "TrimBorder"
+]
+
+
+class ToNumpy:
+
+    def __call__(self, pil_img):
+        np_img = np.array(pil_img, dtype=np.uint8)
+        if np_img.ndim < 3:
+            np_img = np.expand_dims(np_img, axis=-1)
+        np_img = np.rollaxis(np_img, 2)  # HWC to CHW
+        return np_img
+
+
+class ToTensor:
+    """ ToTensor with no rescaling of values"""
+    def __init__(self, dtype=torch.float32):
+        self.dtype = dtype
+
+    def __call__(self, pil_img):
+        return F.pil_to_tensor(pil_img).to(dtype=self.dtype)
+
+
+# Pillow is deprecating the top-level resampling attributes (e.g., Image.BILINEAR) in
+# favor of the Image.Resampling enum. The top-level resampling attributes will be
+# removed in Pillow 10.
+if hasattr(Image, "Resampling"):
+    _pil_interpolation_to_str = {
+        Image.Resampling.NEAREST: 'nearest',
+        Image.Resampling.BILINEAR: 'bilinear',
+        Image.Resampling.BICUBIC: 'bicubic',
+        Image.Resampling.BOX: 'box',
+        Image.Resampling.HAMMING: 'hamming',
+        Image.Resampling.LANCZOS: 'lanczos',
+    }
+else:
+    _pil_interpolation_to_str = {
+        Image.NEAREST: 'nearest',
+        Image.BILINEAR: 'bilinear',
+        Image.BICUBIC: 'bicubic',
+        Image.BOX: 'box',
+        Image.HAMMING: 'hamming',
+        Image.LANCZOS: 'lanczos',
+    }
+
+_str_to_pil_interpolation = {b: a for a, b in _pil_interpolation_to_str.items()}
+
+
+if has_interpolation_mode:
+    _torch_interpolation_to_str = {
+        InterpolationMode.NEAREST: 'nearest',
+        InterpolationMode.BILINEAR: 'bilinear',
+        InterpolationMode.BICUBIC: 'bicubic',
+        InterpolationMode.BOX: 'box',
+        InterpolationMode.HAMMING: 'hamming',
+        InterpolationMode.LANCZOS: 'lanczos',
+    }
+    _str_to_torch_interpolation = {b: a for a, b in _torch_interpolation_to_str.items()}
+else:
+    _pil_interpolation_to_torch = {}
+    _torch_interpolation_to_str = {}
+
+
+def str_to_pil_interp(mode_str):
+    return _str_to_pil_interpolation[mode_str]
+
+
+def str_to_interp_mode(mode_str):
+    if has_interpolation_mode:
+        return _str_to_torch_interpolation[mode_str]
+    else:
+        return _str_to_pil_interpolation[mode_str]
+
+
+def interp_mode_to_str(mode):
+    if has_interpolation_mode:
+        return _torch_interpolation_to_str[mode]
+    else:
+        return _pil_interpolation_to_str[mode]
+
+
+_RANDOM_INTERPOLATION = (str_to_interp_mode('bilinear'), str_to_interp_mode('bicubic'))
+
+
+def _setup_size(size, error_msg="Please provide only two dimensions (h, w) for size."):
+    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.
+
+    A crop of random size (default: of 0.08 to 1.0) of the original size and a random
+    aspect ratio (default: of 3/4 to 4/3) of the original aspect ratio is made. This crop
+    is finally resized to given size.
+    This is popularly used to train the Inception networks.
+
+    Args:
+        size: expected output size of each edge
+        scale: range of size of the origin size cropped
+        ratio: range of aspect ratio of the origin aspect ratio cropped
+        interpolation: Default: PIL.Image.BILINEAR
+    """
+
+    def __init__(
+            self,
+            size,
+            scale=(0.08, 1.0),
+            ratio=(3. / 4., 4. / 3.),
+            interpolation='bilinear',
+    ):
+        if isinstance(size, (list, tuple)):
+            self.size = tuple(size)
+        else:
+            self.size = (size, size)
+        if (scale[0] > scale[1]) or (ratio[0] > ratio[1]):
+            warnings.warn("range should be of kind (min, max)")
+
+        if interpolation == 'random':
+            self.interpolation = _RANDOM_INTERPOLATION
+        else:
+            self.interpolation = str_to_interp_mode(interpolation)
+        self.scale = scale
+        self.ratio = ratio
+
+    @staticmethod
+    def get_params(img, scale, ratio):
+        """Get parameters for ``crop`` for a random sized crop.
+
+        Args:
+            img (PIL Image): Image to be cropped.
+            scale (tuple): range of size of the origin size cropped
+            ratio (tuple): range of aspect ratio of the origin aspect ratio cropped
+
+        Returns:
+            tuple: params (i, j, h, w) to be passed to ``crop`` for a random
+                sized crop.
+        """
+        img_w, img_h = F.get_image_size(img)
+        area = img_w * img_h
+
+        for attempt in range(10):
+            target_area = random.uniform(*scale) * area
+            log_ratio = (math.log(ratio[0]), math.log(ratio[1]))
+            aspect_ratio = math.exp(random.uniform(*log_ratio))
+
+            target_w = int(round(math.sqrt(target_area * aspect_ratio)))
+            target_h = int(round(math.sqrt(target_area / aspect_ratio)))
+            if target_w <= img_w and target_h <= img_h:
+                i = random.randint(0, img_h - target_h)
+                j = random.randint(0, img_w - target_w)
+                return i, j, target_h, target_w
+
+        # Fallback to central crop
+        in_ratio = img_w / img_h
+        if in_ratio < min(ratio):
+            target_w = img_w
+            target_h = int(round(target_w / min(ratio)))
+        elif in_ratio > max(ratio):
+            target_h = img_h
+            target_w = int(round(target_h * max(ratio)))
+        else:  # whole image
+            target_w = img_w
+            target_h = img_h
+        i = (img_h - target_h) // 2
+        j = (img_w - target_w) // 2
+        return i, j, target_h, target_w
+
+    def __call__(self, img):
+        """
+        Args:
+            img (PIL Image): Image to be cropped and resized.
+
+        Returns:
+            PIL Image: Randomly cropped and resized image.
+        """
+        i, j, h, w = self.get_params(img, self.scale, self.ratio)
+        if isinstance(self.interpolation, (tuple, list)):
+            interpolation = random.choice(self.interpolation)
+        else:
+            interpolation = self.interpolation
+        return F.resized_crop(img, i, j, h, w, self.size, interpolation)
+
+    def __repr__(self):
+        if isinstance(self.interpolation, (tuple, list)):
+            interpolate_str = ' '.join([interp_mode_to_str(x) for x in self.interpolation])
+        else:
+            interpolate_str = interp_mode_to_str(self.interpolation)
+        format_string = self.__class__.__name__ + '(size={0}'.format(self.size)
+        format_string += ', scale={0}'.format(tuple(round(s, 4) for s in self.scale))
+        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: Union[int, List[int]],
+        fill: Union[int, Tuple[int, int, int]] = 0,
+        padding_mode: str = 'constant',
+) -> 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.
+    """
+    output_size = _setup_size(output_size)
+    crop_height, crop_width = output_size
+    _, image_height, image_width = F.get_dimensions(img)
+
+    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, padding_mode=padding_mode)
+        _, 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: Union[int, List[int]],
+            fill: Union[int, Tuple[int, int, int]] = 0,
+            padding_mode: str = 'constant',
+    ):
+        super().__init__()
+        self.size = _setup_size(size)
+        self.fill = fill
+        self.padding_mode = padding_mode
+
+    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, padding_mode=self.padding_mode)
+
+    def __repr__(self) -> str:
+        return f"{self.__class__.__name__}(size={self.size})"
+
+
+def crop_or_pad(
+        img: torch.Tensor,
+        top: int,
+        left: int,
+        height: int,
+        width: int,
+        fill: Union[int, Tuple[int, int, int]] = 0,
+        padding_mode: str = 'constant',
+) -> torch.Tensor:
+    """ Crops and/or pads image to meet target size, with control over fill and padding_mode.
+    """
+    _, image_height, image_width = F.get_dimensions(img)
+    right = left + width
+    bottom = top + height
+    if left < 0 or top < 0 or right > image_width or bottom > image_height:
+        padding_ltrb = [
+            max(-left + min(0, right), 0),
+            max(-top + min(0, bottom), 0),
+            max(right - max(image_width, left), 0),
+            max(bottom - max(image_height, top), 0),
+        ]
+        img = F.pad(img, padding_ltrb, fill=fill, padding_mode=padding_mode)
+
+    top = max(top, 0)
+    left = max(left, 0)
+    return F.crop(img, top, left, height, width)
+
+
+class RandomCropOrPad(torch.nn.Module):
+    """ Crop and/or pad image with random placement within the crop or pad margin.
+    """
+
+    def __init__(
+            self,
+            size: Union[int, List[int]],
+            fill: Union[int, Tuple[int, int, int]] = 0,
+            padding_mode: str = 'constant',
+    ):
+        super().__init__()
+        self.size = _setup_size(size)
+        self.fill = fill
+        self.padding_mode = padding_mode
+
+    @staticmethod
+    def get_params(img, size):
+        _, image_height, image_width = F.get_dimensions(img)
+        delta_height = image_height - size[0]
+        delta_width = image_width - size[1]
+        top = int(math.copysign(random.randint(0, abs(delta_height)), delta_height))
+        left = int(math.copysign(random.randint(0, abs(delta_width)), delta_width))
+        return top, left
+
+    def forward(self, img):
+        """
+        Args:
+            img (PIL Image or Tensor): Image to be cropped.
+
+        Returns:
+            PIL Image or Tensor: Cropped image.
+        """
+        top, left = self.get_params(img, self.size)
+        return crop_or_pad(
+            img,
+            top=top,
+            left=left,
+            height=self.size[0],
+            width=self.size[1],
+            fill=self.fill,
+            padding_mode=self.padding_mode,
+        )
+
+    def __repr__(self) -> str:
+        return f"{self.__class__.__name__}(size={self.size})"
+
+
+class RandomPad:
+    def __init__(self, input_size, fill=0):
+        self.input_size = input_size
+        self.fill = fill
+
+    @staticmethod
+    def get_params(img, input_size):
+        width, height = F.get_image_size(img)
+        delta_width = max(input_size[1] - width, 0)
+        delta_height = max(input_size[0] - height, 0)
+        pad_left = random.randint(0, delta_width)
+        pad_top = random.randint(0, delta_height)
+        pad_right = delta_width - pad_left
+        pad_bottom = delta_height - pad_top
+        return pad_left, pad_top, pad_right, pad_bottom
+
+    def __call__(self, img):
+        padding = self.get_params(img, self.input_size)
+        img = F.pad(img, padding, self.fill)
+        return img
+
+
+class ResizeKeepRatio:
+    """ Resize and Keep Aspect Ratio
+    """
+
+    def __init__(
+            self,
+            size,
+            longest=0.,
+            interpolation='bilinear',
+            random_scale_prob=0.,
+            random_scale_range=(0.85, 1.05),
+            random_scale_area=False,
+            random_aspect_prob=0.,
+            random_aspect_range=(0.9, 1.11),
+    ):
+        """
+
+        Args:
+            size:
+            longest:
+            interpolation:
+            random_scale_prob:
+            random_scale_range:
+            random_scale_area:
+            random_aspect_prob:
+            random_aspect_range:
+        """
+        if isinstance(size, (list, tuple)):
+            self.size = tuple(size)
+        else:
+            self.size = (size, size)
+        if interpolation == 'random':
+            self.interpolation = _RANDOM_INTERPOLATION
+        else:
+            self.interpolation = str_to_interp_mode(interpolation)
+        self.longest = float(longest)
+        self.random_scale_prob = random_scale_prob
+        self.random_scale_range = random_scale_range
+        self.random_scale_area = random_scale_area
+        self.random_aspect_prob = random_aspect_prob
+        self.random_aspect_range = random_aspect_range
+
+    @staticmethod
+    def get_params(
+            img,
+            target_size,
+            longest,
+            random_scale_prob=0.,
+            random_scale_range=(1.0, 1.33),
+            random_scale_area=False,
+            random_aspect_prob=0.,
+            random_aspect_range=(0.9, 1.11)
+    ):
+        """Get parameters
+        """
+        img_h, img_w = img_size = F.get_dimensions(img)[1:]
+        target_h, target_w = target_size
+        ratio_h = img_h / target_h
+        ratio_w = img_w / target_w
+        ratio = max(ratio_h, ratio_w) * longest + min(ratio_h, ratio_w) * (1. - longest)
+
+        if random_scale_prob > 0 and random.random() < random_scale_prob:
+            ratio_factor = random.uniform(random_scale_range[0], random_scale_range[1])
+            if random_scale_area:
+                # make ratio factor equivalent to RRC area crop where < 1.0 = area zoom,
+                # otherwise like affine scale where < 1.0 = linear zoom out
+                ratio_factor = 1. / math.sqrt(ratio_factor)
+            ratio_factor = (ratio_factor, ratio_factor)
+        else:
+            ratio_factor = (1., 1.)
+
+        if random_aspect_prob > 0 and random.random() < random_aspect_prob:
+            log_aspect = (math.log(random_aspect_range[0]), math.log(random_aspect_range[1]))
+            aspect_factor = math.exp(random.uniform(*log_aspect))
+            aspect_factor = math.sqrt(aspect_factor)
+            # currently applying random aspect adjustment equally to both dims,
+            # could change to keep output sizes above their target where possible
+            ratio_factor = (ratio_factor[0] / aspect_factor, ratio_factor[1] * aspect_factor)
+
+        size = [round(x * f / ratio) for x, f in zip(img_size, ratio_factor)]
+        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,
+            self.random_scale_prob, self.random_scale_range, self.random_scale_area,
+            self.random_aspect_prob, self.random_aspect_range
+        )
+        if isinstance(self.interpolation, (tuple, list)):
+            interpolation = random.choice(self.interpolation)
+        else:
+            interpolation = self.interpolation
+        img = F.resize(img, size, interpolation)
+        return img
+
+    def __repr__(self):
+        if isinstance(self.interpolation, (tuple, list)):
+            interpolate_str = ' '.join([interp_mode_to_str(x) for x in self.interpolation])
+        else:
+            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}'
+        format_string += f', random_scale_prob={self.random_scale_prob:.3f}'
+        format_string += f', random_scale_range=(' \
+                         f'{self.random_scale_range[0]:.3f}, {self.random_aspect_range[1]:.3f})'
+        format_string += f', random_aspect_prob={self.random_aspect_prob:.3f}'
+        format_string += f', random_aspect_range=(' \
+                         f'{self.random_aspect_range[0]:.3f}, {self.random_aspect_range[1]:.3f}))'
+        return format_string
+
+
+class TrimBorder(torch.nn.Module):
+
+    def __init__(
+            self,
+            border_size: int,
+    ):
+        super().__init__()
+        self.border_size = border_size
+
+    def forward(self, img):
+        w, h = F.get_image_size(img)
+        top = left = self.border_size
+        top = min(top, h)
+        left = min(left, h)
+        height = max(0, h - 2 * self.border_size)
+        width = max(0, w - 2 * self.border_size)
+        return F.crop(img, top, left, height, width)

timm/data/transforms_factory.py

@@ -0,0 +1,440 @@
+""" Transforms Factory
+Factory methods for building image transforms for use with TIMM (PyTorch Image Models)
+
+Hacked together by / Copyright 2019, Ross Wightman
+"""
+import math
+from typing import Optional, Tuple, Union
+
+import torch
+from torchvision import transforms
+
+from timm.data.constants import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD, DEFAULT_CROP_PCT
+from timm.data.auto_augment import rand_augment_transform, augment_and_mix_transform, auto_augment_transform
+from timm.data.transforms import str_to_interp_mode, str_to_pil_interp, RandomResizedCropAndInterpolation,\
+    ResizeKeepRatio, CenterCropOrPad, RandomCropOrPad, TrimBorder, ToNumpy
+from timm.data.random_erasing import RandomErasing
+
+
+def transforms_noaug_train(
+        img_size: Union[int, Tuple[int, int]] = 224,
+        interpolation: str = 'bilinear',
+        use_prefetcher: bool = False,
+        mean: Tuple[float, ...] = IMAGENET_DEFAULT_MEAN,
+        std: Tuple[float, ...] = IMAGENET_DEFAULT_STD,
+):
+    """ No-augmentation image transforms for training.
+
+    Args:
+        img_size: Target image size.
+        interpolation: Image interpolation mode.
+        mean: Image normalization mean.
+        std: Image normalization standard deviation.
+        use_prefetcher: Prefetcher enabled. Do not convert image to tensor or normalize.
+
+    Returns:
+
+    """
+    if interpolation == 'random':
+        # random interpolation not supported with no-aug
+        interpolation = 'bilinear'
+    tfl = [
+        transforms.Resize(img_size, interpolation=str_to_interp_mode(interpolation)),
+        transforms.CenterCrop(img_size)
+    ]
+    if use_prefetcher:
+        # prefetcher and collate will handle tensor conversion and norm
+        tfl += [ToNumpy()]
+    else:
+        tfl += [
+            transforms.ToTensor(),
+            transforms.Normalize(
+                mean=torch.tensor(mean),
+                std=torch.tensor(std)
+            )
+        ]
+    return transforms.Compose(tfl)
+
+
+def transforms_imagenet_train(
+        img_size: Union[int, Tuple[int, int]] = 224,
+        scale: Optional[Tuple[float, float]] = None,
+        ratio: Optional[Tuple[float, float]] = None,
+        train_crop_mode: Optional[str] = None,
+        hflip: float = 0.5,
+        vflip: float = 0.,
+        color_jitter: Union[float, Tuple[float, ...]] = 0.4,
+        color_jitter_prob: Optional[float] = None,
+        force_color_jitter: bool = False,
+        grayscale_prob: float = 0.,
+        gaussian_blur_prob: float = 0.,
+        auto_augment: Optional[str] = None,
+        interpolation: str = 'random',
+        mean: Tuple[float, ...] = IMAGENET_DEFAULT_MEAN,
+        std: Tuple[float, ...] = IMAGENET_DEFAULT_STD,
+        re_prob: float = 0.,
+        re_mode: str = 'const',
+        re_count: int = 1,
+        re_num_splits: int = 0,
+        use_prefetcher: bool = False,
+        separate: bool = False,
+):
+    """ ImageNet-oriented image transforms for training.
+
+    Args:
+        img_size: Target image size.
+        train_crop_mode: Training random crop mode ('rrc', 'rkrc', 'rkrr').
+        scale: Random resize scale range (crop area, < 1.0 => zoom in).
+        ratio: Random aspect ratio range (crop ratio for RRC, ratio adjustment factor for RKR).
+        hflip: Horizontal flip probability.
+        vflip: Vertical flip probability.
+        color_jitter: Random color jitter component factors (brightness, contrast, saturation, hue).
+            Scalar is applied as (scalar,) * 3 (no hue).
+        color_jitter_prob: Apply color jitter with this probability if not None (for SimlCLR-like aug).
+        force_color_jitter: Force color jitter where it is normally disabled (ie with RandAugment on).
+        grayscale_prob: Probability of converting image to grayscale (for SimCLR-like aug).
+        gaussian_blur_prob: Probability of applying gaussian blur (for SimCLR-like aug).
+        auto_augment: Auto augment configuration string (see auto_augment.py).
+        interpolation: Image interpolation mode.
+        mean: Image normalization mean.
+        std: Image normalization standard deviation.
+        re_prob: Random erasing probability.
+        re_mode: Random erasing fill mode.
+        re_count: Number of random erasing regions.
+        re_num_splits: Control split of random erasing across batch size.
+        use_prefetcher: Prefetcher enabled. Do not convert image to tensor or normalize.
+        separate: Output transforms in 3-stage tuple.
+
+    Returns:
+        If separate==True, the transforms are returned as a tuple of 3 separate transforms
+        for use in a mixing dataset that passes
+         * all data through the first (primary) transform, called the 'clean' data
+         * a portion of the data through the secondary transform
+         * normalizes and converts the branches above with the third, final transform
+    """
+    train_crop_mode = train_crop_mode or 'rrc'
+    assert train_crop_mode in {'rrc', 'rkrc', 'rkrr'}
+    if train_crop_mode in ('rkrc', 'rkrr'):
+        # FIXME integration of RKR is a WIP
+        scale = tuple(scale or (0.8, 1.00))
+        ratio = tuple(ratio or (0.9, 1/.9))
+        primary_tfl = [
+            ResizeKeepRatio(
+                img_size,
+                interpolation=interpolation,
+                random_scale_prob=0.5,
+                random_scale_range=scale,
+                random_scale_area=True,  # scale compatible with RRC
+                random_aspect_prob=0.5,
+                random_aspect_range=ratio,
+            ),
+            CenterCropOrPad(img_size, padding_mode='reflect')
+            if train_crop_mode == 'rkrc' else
+            RandomCropOrPad(img_size, padding_mode='reflect')
+        ]
+    else:
+        scale = tuple(scale or (0.08, 1.0))  # default imagenet scale range
+        ratio = tuple(ratio or (3. / 4., 4. / 3.))  # default imagenet ratio range
+        primary_tfl = [
+            RandomResizedCropAndInterpolation(
+                img_size,
+                scale=scale,
+                ratio=ratio,
+                interpolation=interpolation,
+            )
+        ]
+    if hflip > 0.:
+        primary_tfl += [transforms.RandomHorizontalFlip(p=hflip)]
+    if vflip > 0.:
+        primary_tfl += [transforms.RandomVerticalFlip(p=vflip)]
+
+    secondary_tfl = []
+    disable_color_jitter = False
+    if auto_augment:
+        assert isinstance(auto_augment, str)
+        # color jitter is typically disabled if AA/RA on,
+        # this allows override without breaking old hparm cfgs
+        disable_color_jitter = not (force_color_jitter or '3a' in auto_augment)
+        if isinstance(img_size, (tuple, list)):
+            img_size_min = min(img_size)
+        else:
+            img_size_min = img_size
+        aa_params = dict(
+            translate_const=int(img_size_min * 0.45),
+            img_mean=tuple([min(255, round(255 * x)) for x in mean]),
+        )
+        if interpolation and interpolation != 'random':
+            aa_params['interpolation'] = str_to_pil_interp(interpolation)
+        if auto_augment.startswith('rand'):
+            secondary_tfl += [rand_augment_transform(auto_augment, aa_params)]
+        elif auto_augment.startswith('augmix'):
+            aa_params['translate_pct'] = 0.3
+            secondary_tfl += [augment_and_mix_transform(auto_augment, aa_params)]
+        else:
+            secondary_tfl += [auto_augment_transform(auto_augment, aa_params)]
+
+    if color_jitter is not None and not disable_color_jitter:
+        # color jitter is enabled when not using AA or when forced
+        if isinstance(color_jitter, (list, tuple)):
+            # color jitter should be a 3-tuple/list if spec brightness/contrast/saturation
+            # or 4 if also augmenting hue
+            assert len(color_jitter) in (3, 4)
+        else:
+            # if it's a scalar, duplicate for brightness, contrast, and saturation, no hue
+            color_jitter = (float(color_jitter),) * 3
+        if color_jitter_prob is not None:
+            secondary_tfl += [
+                transforms.RandomApply([
+                        transforms.ColorJitter(*color_jitter),
+                    ],
+                    p=color_jitter_prob
+                )
+            ]
+        else:
+            secondary_tfl += [transforms.ColorJitter(*color_jitter)]
+
+    if grayscale_prob:
+        secondary_tfl += [transforms.RandomGrayscale(p=grayscale_prob)]
+
+    if gaussian_blur_prob:
+        secondary_tfl += [
+            transforms.RandomApply([
+                    transforms.GaussianBlur(kernel_size=23),  # hardcoded for now
+                ],
+                p=gaussian_blur_prob,
+            )
+        ]
+
+    final_tfl = []
+    if use_prefetcher:
+        # prefetcher and collate will handle tensor conversion and norm
+        final_tfl += [ToNumpy()]
+    else:
+        final_tfl += [
+            transforms.ToTensor(),
+            transforms.Normalize(
+                mean=torch.tensor(mean),
+                std=torch.tensor(std)
+            ),
+        ]
+        if re_prob > 0.:
+            final_tfl += [
+                RandomErasing(
+                    re_prob,
+                    mode=re_mode,
+                    max_count=re_count,
+                    num_splits=re_num_splits,
+                    device='cpu',
+                )
+            ]
+
+    if separate:
+        return transforms.Compose(primary_tfl), transforms.Compose(secondary_tfl), transforms.Compose(final_tfl)
+    else:
+        return transforms.Compose(primary_tfl + secondary_tfl + final_tfl)
+
+
+def transforms_imagenet_eval(
+        img_size: Union[int, Tuple[int, int]] = 224,
+        crop_pct: Optional[float] = None,
+        crop_mode: Optional[str] = None,
+        crop_border_pixels: Optional[int] = None,
+        interpolation: str = 'bilinear',
+        mean: Tuple[float, ...] = IMAGENET_DEFAULT_MEAN,
+        std: Tuple[float, ...] = IMAGENET_DEFAULT_STD,
+        use_prefetcher: bool = False,
+):
+    """ ImageNet-oriented image transform for evaluation and inference.
+
+    Args:
+        img_size: Target image size.
+        crop_pct: Crop percentage. Defaults to 0.875 when None.
+        crop_mode: Crop mode. One of ['squash', 'border', 'center']. Defaults to 'center' when None.
+        crop_border_pixels: Trim a border of specified # pixels around edge of original image.
+        interpolation: Image interpolation mode.
+        mean: Image normalization mean.
+        std: Image normalization standard deviation.
+        use_prefetcher: Prefetcher enabled. Do not convert image to tensor or normalize.
+
+    Returns:
+        Composed transform pipeline
+    """
+    crop_pct = crop_pct or DEFAULT_CROP_PCT
+
+    if isinstance(img_size, (tuple, list)):
+        assert len(img_size) == 2
+        scale_size = tuple([math.floor(x / crop_pct) for x in img_size])
+    else:
+        scale_size = math.floor(img_size / crop_pct)
+        scale_size = (scale_size, scale_size)
+
+    tfl = []
+
+    if crop_border_pixels:
+        tfl += [TrimBorder(crop_border_pixels)]
+
+    if crop_mode == 'squash':
+        # squash mode scales each edge to 1/pct of target, then crops
+        # aspect ratio is not preserved, no img lost if crop_pct == 1.0
+        tfl += [
+            transforms.Resize(scale_size, interpolation=str_to_interp_mode(interpolation)),
+            transforms.CenterCrop(img_size),
+        ]
+    elif crop_mode == 'border':
+        # scale the longest edge of image to 1/pct of target edge, add borders to pad, then crop
+        # no image lost if crop_pct == 1.0
+        fill = [round(255 * v) for v in mean]
+        tfl += [
+            ResizeKeepRatio(scale_size, interpolation=interpolation, longest=1.0),
+            CenterCropOrPad(img_size, fill=fill),
+        ]
+    else:
+        # default crop model is center
+        # aspect ratio is preserved, crops center within image, no borders are added, image is lost
+        if scale_size[0] == scale_size[1]:
+            # simple case, use torchvision built-in Resize w/ shortest edge mode (scalar size arg)
+            tfl += [
+                transforms.Resize(scale_size[0], interpolation=str_to_interp_mode(interpolation))
+            ]
+        else:
+            # resize the shortest edge to matching target dim for non-square target
+            tfl += [ResizeKeepRatio(scale_size)]
+        tfl += [transforms.CenterCrop(img_size)]
+
+    if use_prefetcher:
+        # prefetcher and collate will handle tensor conversion and norm
+        tfl += [ToNumpy()]
+    else:
+        tfl += [
+            transforms.ToTensor(),
+            transforms.Normalize(
+                mean=torch.tensor(mean),
+                std=torch.tensor(std),
+            )
+        ]
+
+    return transforms.Compose(tfl)
+
+
+def create_transform(
+        input_size: Union[int, Tuple[int, int], Tuple[int, int, int]] = 224,
+        is_training: bool = False,
+        no_aug: bool = False,
+        train_crop_mode: Optional[str] = None,
+        scale: Optional[Tuple[float, float]] = None,
+        ratio: Optional[Tuple[float, float]] = None,
+        hflip: float = 0.5,
+        vflip: float = 0.,
+        color_jitter: Union[float, Tuple[float, ...]] = 0.4,
+        color_jitter_prob: Optional[float] = None,
+        grayscale_prob: float = 0.,
+        gaussian_blur_prob: float = 0.,
+        auto_augment: Optional[str] = None,
+        interpolation: str = 'bilinear',
+        mean: Tuple[float, ...] = IMAGENET_DEFAULT_MEAN,
+        std: Tuple[float, ...] = IMAGENET_DEFAULT_STD,
+        re_prob: float = 0.,
+        re_mode: str = 'const',
+        re_count: int = 1,
+        re_num_splits: int = 0,
+        crop_pct: Optional[float] = None,
+        crop_mode: Optional[str] = None,
+        crop_border_pixels: Optional[int] = None,
+        tf_preprocessing: bool = False,
+        use_prefetcher: bool = False,
+        separate: bool = False,
+):
+    """
+
+    Args:
+        input_size: Target input size (channels, height, width) tuple or size scalar.
+        is_training: Return training (random) transforms.
+        no_aug: Disable augmentation for training (useful for debug).
+        train_crop_mode: Training random crop mode ('rrc', 'rkrc', 'rkrr').
+        scale: Random resize scale range (crop area, < 1.0 => zoom in).
+        ratio: Random aspect ratio range (crop ratio for RRC, ratio adjustment factor for RKR).
+        hflip: Horizontal flip probability.
+        vflip: Vertical flip probability.
+        color_jitter: Random color jitter component factors (brightness, contrast, saturation, hue).
+            Scalar is applied as (scalar,) * 3 (no hue).
+        color_jitter_prob: Apply color jitter with this probability if not None (for SimlCLR-like aug).
+        grayscale_prob: Probability of converting image to grayscale (for SimCLR-like aug).
+        gaussian_blur_prob: Probability of applying gaussian blur (for SimCLR-like aug).
+        auto_augment: Auto augment configuration string (see auto_augment.py).
+        interpolation: Image interpolation mode.
+        mean: Image normalization mean.
+        std: Image normalization standard deviation.
+        re_prob: Random erasing probability.
+        re_mode: Random erasing fill mode.
+        re_count: Number of random erasing regions.
+        re_num_splits: Control split of random erasing across batch size.
+        crop_pct: Inference crop percentage (output size / resize size).
+        crop_mode: Inference crop mode. One of ['squash', 'border', 'center']. Defaults to 'center' when None.
+        crop_border_pixels: Inference crop border of specified # pixels around edge of original image.
+        tf_preprocessing: Use TF 1.0 inference preprocessing for testing model ports
+        use_prefetcher: Pre-fetcher enabled. Do not convert image to tensor or normalize.
+        separate: Output transforms in 3-stage tuple.
+
+    Returns:
+        Composed transforms or tuple thereof
+    """
+    if isinstance(input_size, (tuple, list)):
+        img_size = input_size[-2:]
+    else:
+        img_size = input_size
+
+    if tf_preprocessing and use_prefetcher:
+        assert not separate, "Separate transforms not supported for TF preprocessing"
+        from timm.data.tf_preprocessing import TfPreprocessTransform
+        transform = TfPreprocessTransform(
+            is_training=is_training,
+            size=img_size,
+            interpolation=interpolation,
+        )
+    else:
+        if is_training and no_aug:
+            assert not separate, "Cannot perform split augmentation with no_aug"
+            transform = transforms_noaug_train(
+                img_size,
+                interpolation=interpolation,
+                use_prefetcher=use_prefetcher,
+                mean=mean,
+                std=std,
+            )
+        elif is_training:
+            transform = transforms_imagenet_train(
+                img_size,
+                train_crop_mode=train_crop_mode,
+                scale=scale,
+                ratio=ratio,
+                hflip=hflip,
+                vflip=vflip,
+                color_jitter=color_jitter,
+                color_jitter_prob=color_jitter_prob,
+                grayscale_prob=grayscale_prob,
+                gaussian_blur_prob=gaussian_blur_prob,
+                auto_augment=auto_augment,
+                interpolation=interpolation,
+                use_prefetcher=use_prefetcher,
+                mean=mean,
+                std=std,
+                re_prob=re_prob,
+                re_mode=re_mode,
+                re_count=re_count,
+                re_num_splits=re_num_splits,
+                separate=separate,
+            )
+        else:
+            assert not separate, "Separate transforms not supported for validation preprocessing"
+            transform = transforms_imagenet_eval(
+                img_size,
+                interpolation=interpolation,
+                use_prefetcher=use_prefetcher,
+                mean=mean,
+                std=std,
+                crop_pct=crop_pct,
+                crop_mode=crop_mode,
+                crop_border_pixels=crop_border_pixels,
+            )
+
+    return transform

timm/layers/__init__.py

@@ -0,0 +1,57 @@
+from .activations import *
+from .adaptive_avgmax_pool import \
+    adaptive_avgmax_pool2d, select_adaptive_pool2d, AdaptiveAvgMaxPool2d, SelectAdaptivePool2d
+from .attention_pool import AttentionPoolLatent
+from .attention_pool2d import AttentionPool2d, RotAttentionPool2d, RotaryEmbedding
+from .blur_pool import BlurPool2d
+from .classifier import ClassifierHead, create_classifier, NormMlpClassifierHead
+from .cond_conv2d import CondConv2d, get_condconv_initializer
+from .config import is_exportable, is_scriptable, is_no_jit, use_fused_attn, \
+    set_exportable, set_scriptable, set_no_jit, set_layer_config, set_fused_attn
+from .conv2d_same import Conv2dSame, conv2d_same
+from .conv_bn_act import ConvNormAct, ConvNormActAa, ConvBnAct
+from .create_act import create_act_layer, get_act_layer, get_act_fn
+from .create_attn import get_attn, create_attn
+from .create_conv2d import create_conv2d
+from .create_norm import get_norm_layer, create_norm_layer
+from .create_norm_act import get_norm_act_layer, create_norm_act_layer, get_norm_act_layer
+from .drop import DropBlock2d, DropPath, drop_block_2d, drop_path
+from .eca import EcaModule, CecaModule, EfficientChannelAttn, CircularEfficientChannelAttn
+from .evo_norm import EvoNorm2dB0, EvoNorm2dB1, EvoNorm2dB2,\
+    EvoNorm2dS0, EvoNorm2dS0a, EvoNorm2dS1, EvoNorm2dS1a, EvoNorm2dS2, EvoNorm2dS2a
+from .fast_norm import is_fast_norm, set_fast_norm, fast_group_norm, fast_layer_norm
+from .filter_response_norm import FilterResponseNormTlu2d, FilterResponseNormAct2d
+from .format import Format, get_channel_dim, get_spatial_dim, nchw_to, nhwc_to
+from .gather_excite import GatherExcite
+from .global_context import GlobalContext
+from .grid import ndgrid, meshgrid
+from .helpers import to_ntuple, to_2tuple, to_3tuple, to_4tuple, make_divisible, extend_tuple
+from .inplace_abn import InplaceAbn
+from .linear import Linear
+from .mixed_conv2d import MixedConv2d
+from .mlp import Mlp, GluMlp, GatedMlp, SwiGLU, SwiGLUPacked, ConvMlp, GlobalResponseNormMlp
+from .non_local_attn import NonLocalAttn, BatNonLocalAttn
+from .norm import GroupNorm, GroupNorm1, LayerNorm, LayerNorm2d, RmsNorm
+from .norm_act import BatchNormAct2d, GroupNormAct, GroupNorm1Act, LayerNormAct, LayerNormAct2d,\
+    SyncBatchNormAct, convert_sync_batchnorm, FrozenBatchNormAct2d, freeze_batch_norm_2d, unfreeze_batch_norm_2d
+from .padding import get_padding, get_same_padding, pad_same
+from .patch_dropout import PatchDropout
+from .patch_embed import PatchEmbed, PatchEmbedWithSize, resample_patch_embed
+from .pool2d_same import AvgPool2dSame, create_pool2d
+from .pos_embed import resample_abs_pos_embed, resample_abs_pos_embed_nhwc
+from .pos_embed_rel import RelPosMlp, RelPosBias, RelPosBiasTf, gen_relative_position_index, gen_relative_log_coords, \
+    resize_rel_pos_bias_table, resize_rel_pos_bias_table_simple, resize_rel_pos_bias_table_levit
+from .pos_embed_sincos import pixel_freq_bands, freq_bands, build_sincos2d_pos_embed, build_fourier_pos_embed, \
+    build_rotary_pos_embed, apply_rot_embed, apply_rot_embed_cat, apply_rot_embed_list, apply_keep_indices_nlc, \
+    FourierEmbed, RotaryEmbedding, RotaryEmbeddingCat
+from .squeeze_excite import SEModule, SqueezeExcite, EffectiveSEModule, EffectiveSqueezeExcite
+from .selective_kernel import SelectiveKernel
+from .separable_conv import SeparableConv2d, SeparableConvNormAct
+from .space_to_depth import SpaceToDepthModule, SpaceToDepth, DepthToSpace
+from .split_attn import SplitAttn
+from .split_batchnorm import SplitBatchNorm2d, convert_splitbn_model
+from .std_conv import StdConv2d, StdConv2dSame, ScaledStdConv2d, ScaledStdConv2dSame
+from .test_time_pool import TestTimePoolHead, apply_test_time_pool
+from .trace_utils import _assert, _float_to_int
+from .typing import LayerType, PadType
+from .weight_init import trunc_normal_, trunc_normal_tf_, variance_scaling_, lecun_normal_

timm/layers/activations.py

@@ -0,0 +1,173 @@
+""" Activations
+
+A collection of activations fn and modules with a common interface so that they can
+easily be swapped. All have an `inplace` arg even if not used.
+
+Hacked together by / Copyright 2020 Ross Wightman
+"""
+
+import torch
+from torch import nn as nn
+from torch.nn import functional as F
+
+
+def swish(x, inplace: bool = False):
+    """Swish - Described in: https://arxiv.org/abs/1710.05941
+    """
+    return x.mul_(x.sigmoid()) if inplace else x.mul(x.sigmoid())
+
+
+class Swish(nn.Module):
+    def __init__(self, inplace: bool = False):
+        super(Swish, self).__init__()
+        self.inplace = inplace
+
+    def forward(self, x):
+        return swish(x, self.inplace)
+
+
+def mish(x, inplace: bool = False):
+    """Mish: A Self Regularized Non-Monotonic Neural Activation Function - https://arxiv.org/abs/1908.08681
+    NOTE: I don't have a working inplace variant
+    """
+    return x.mul(F.softplus(x).tanh())
+
+
+class Mish(nn.Module):
+    """Mish: A Self Regularized Non-Monotonic Neural Activation Function - https://arxiv.org/abs/1908.08681
+    """
+    def __init__(self, inplace: bool = False):
+        super(Mish, self).__init__()
+
+    def forward(self, x):
+        return mish(x)
+
+
+def sigmoid(x, inplace: bool = False):
+    return x.sigmoid_() if inplace else x.sigmoid()
+
+
+# PyTorch has this, but not with a consistent inplace argmument interface
+class Sigmoid(nn.Module):
+    def __init__(self, inplace: bool = False):
+        super(Sigmoid, self).__init__()
+        self.inplace = inplace
+
+    def forward(self, x):
+        return x.sigmoid_() if self.inplace else x.sigmoid()
+
+
+def tanh(x, inplace: bool = False):
+    return x.tanh_() if inplace else x.tanh()
+
+
+# PyTorch has this, but not with a consistent inplace argmument interface
+class Tanh(nn.Module):
+    def __init__(self, inplace: bool = False):
+        super(Tanh, self).__init__()
+        self.inplace = inplace
+
+    def forward(self, x):
+        return x.tanh_() if self.inplace else x.tanh()
+
+
+def hard_swish(x, inplace: bool = False):
+    inner = F.relu6(x + 3.).div_(6.)
+    return x.mul_(inner) if inplace else x.mul(inner)
+
+
+class HardSwish(nn.Module):
+    def __init__(self, inplace: bool = False):
+        super(HardSwish, self).__init__()
+        self.inplace = inplace
+
+    def forward(self, x):
+        return hard_swish(x, self.inplace)
+
+
+def hard_sigmoid(x, inplace: bool = False):
+    if inplace:
+        return x.add_(3.).clamp_(0., 6.).div_(6.)
+    else:
+        return F.relu6(x + 3.) / 6.
+
+
+class HardSigmoid(nn.Module):
+    def __init__(self, inplace: bool = False):
+        super(HardSigmoid, self).__init__()
+        self.inplace = inplace
+
+    def forward(self, x):
+        return hard_sigmoid(x, self.inplace)
+
+
+def hard_mish(x, inplace: bool = False):
+    """ Hard Mish
+    Experimental, based on notes by Mish author Diganta Misra at
+      https://github.com/digantamisra98/H-Mish/blob/0da20d4bc58e696b6803f2523c58d3c8a82782d0/README.md
+    """
+    if inplace:
+        return x.mul_(0.5 * (x + 2).clamp(min=0, max=2))
+    else:
+        return 0.5 * x * (x + 2).clamp(min=0, max=2)
+
+
+class HardMish(nn.Module):
+    def __init__(self, inplace: bool = False):
+        super(HardMish, self).__init__()
+        self.inplace = inplace
+
+    def forward(self, x):
+        return hard_mish(x, self.inplace)
+
+
+class PReLU(nn.PReLU):
+    """Applies PReLU (w/ dummy inplace arg)
+    """
+    def __init__(self, num_parameters: int = 1, init: float = 0.25, inplace: bool = False) -> None:
+        super(PReLU, self).__init__(num_parameters=num_parameters, init=init)
+
+    def forward(self, input: torch.Tensor) -> torch.Tensor:
+        return F.prelu(input, self.weight)
+
+
+def gelu(x: torch.Tensor, inplace: bool = False) -> torch.Tensor:
+    return F.gelu(x)
+
+
+class GELU(nn.Module):
+    """Applies the Gaussian Error Linear Units function (w/ dummy inplace arg)
+    """
+    def __init__(self, inplace: bool = False):
+        super(GELU, self).__init__()
+
+    def forward(self, input: torch.Tensor) -> torch.Tensor:
+        return F.gelu(input)
+
+
+def gelu_tanh(x: torch.Tensor, inplace: bool = False) -> torch.Tensor:
+    return F.gelu(x, approximate='tanh')
+
+
+class GELUTanh(nn.Module):
+    """Applies the Gaussian Error Linear Units function (w/ dummy inplace arg)
+    """
+    def __init__(self, inplace: bool = False):
+        super(GELUTanh, self).__init__()
+
+    def forward(self, input: torch.Tensor) -> torch.Tensor:
+        return F.gelu(input, approximate='tanh')
+
+
+def quick_gelu(x: torch.Tensor, inplace: bool = False) -> torch.Tensor:
+    return x * torch.sigmoid(1.702 * x)
+
+
+class QuickGELU(nn.Module):
+    """Applies the Gaussian Error Linear Units function (w/ dummy inplace arg)
+    """
+    def __init__(self, inplace: bool = False):
+        super(QuickGELU, self).__init__()
+
+    def forward(self, input: torch.Tensor) -> torch.Tensor:
+        return quick_gelu(input)

timm/layers/activations_jit.py

@@ -0,0 +1,90 @@
+""" Activations
+
+A collection of jit-scripted activations fn and modules with a common interface so that they can
+easily be swapped. All have an `inplace` arg even if not used.
+
+All jit scripted activations are lacking in-place variations on purpose, scripted kernel fusion does not
+currently work across in-place op boundaries, thus performance is equal to or less than the non-scripted
+versions if they contain in-place ops.
+
+Hacked together by / Copyright 2020 Ross Wightman
+"""
+
+import torch
+from torch import nn as nn
+from torch.nn import functional as F
+
+
+@torch.jit.script
+def swish_jit(x, inplace: bool = False):
+    """Swish - Described in: https://arxiv.org/abs/1710.05941
+    """
+    return x.mul(x.sigmoid())
+
+
+@torch.jit.script
+def mish_jit(x, _inplace: bool = False):
+    """Mish: A Self Regularized Non-Monotonic Neural Activation Function - https://arxiv.org/abs/1908.08681
+    """
+    return x.mul(F.softplus(x).tanh())
+
+
+class SwishJit(nn.Module):
+    def __init__(self, inplace: bool = False):
+        super(SwishJit, self).__init__()
+
+    def forward(self, x):
+        return swish_jit(x)
+
+
+class MishJit(nn.Module):
+    def __init__(self, inplace: bool = False):
+        super(MishJit, self).__init__()
+
+    def forward(self, x):
+        return mish_jit(x)
+
+
+@torch.jit.script
+def hard_sigmoid_jit(x, inplace: bool = False):
+    # return F.relu6(x + 3.) / 6.
+    return (x + 3).clamp(min=0, max=6).div(6.)  # clamp seems ever so slightly faster?
+
+
+class HardSigmoidJit(nn.Module):
+    def __init__(self, inplace: bool = False):
+        super(HardSigmoidJit, self).__init__()
+
+    def forward(self, x):
+        return hard_sigmoid_jit(x)
+
+
+@torch.jit.script
+def hard_swish_jit(x, inplace: bool = False):
+    # return x * (F.relu6(x + 3.) / 6)
+    return x * (x + 3).clamp(min=0, max=6).div(6.)  # clamp seems ever so slightly faster?
+
+
+class HardSwishJit(nn.Module):
+    def __init__(self, inplace: bool = False):
+        super(HardSwishJit, self).__init__()
+
+    def forward(self, x):
+        return hard_swish_jit(x)
+
+
+@torch.jit.script
+def hard_mish_jit(x, inplace: bool = False):
+    """ Hard Mish
+    Experimental, based on notes by Mish author Diganta Misra at
+      https://github.com/digantamisra98/H-Mish/blob/0da20d4bc58e696b6803f2523c58d3c8a82782d0/README.md
+    """
+    return 0.5 * x * (x + 2).clamp(min=0, max=2)
+
+
+class HardMishJit(nn.Module):
+    def __init__(self, inplace: bool = False):
+        super(HardMishJit, self).__init__()
+
+    def forward(self, x):
+        return hard_mish_jit(x)