Merge branch 'main' into fast_load

Author: gau-nernst

timm/data/dataset.py

@@ -103,6 +103,7 @@ def __init__(
             transform=None,
             target_transform=None,
             max_steps=None,
+            **kwargs,
     ):
         assert reader is not None
         if isinstance(reader, str):
@@ -121,6 +122,7 @@ def __init__(
                 input_key=input_key,
                 target_key=target_key,
                 max_steps=max_steps,
+                **kwargs,
             )
         else:
             self.reader = reader

timm/data/dataset_factory.py

@@ -74,34 +74,37 @@ def create_dataset(
         seed: int = 42,
         repeats: int = 0,
         input_img_mode: str = 'RGB',
+        trust_remote_code: bool = False,
         **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
+      * Folder - default, timm folder (or tar) based ImageDataset
+      * Torch - torchvision based datasets
       * HFDS - Hugging Face Datasets
+      * HFIDS - Hugging Face Datasets Iterable (streaming mode, with IterableDataset)
       * TFDS - Tensorflow-datasets wrapper in IterabeDataset interface via IterableImageDataset
       * WDS - Webdataset
-      * all - any of the above
+      * 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
+        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)
+        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 (HFIDS, 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, HFIDS)
+        batch_size: Batch size hint for iterable datasets (TFDS, WDS, HFIDS)
+        seed: Seed for iterable datasets (TFDS, WDS, HFIDS)
+        repeats: Dataset repeats per iteration i.e. epoch (TFDS, WDS, HFIDS)
+        input_img_mode: Input image color conversion mode e.g. 'RGB', 'L' (folder, TFDS, WDS, HFDS, HFIDS)
+        trust_remote_code: Trust remote code in Hugging Face Datasets if True (HFDS, HFIDS)
+        **kwargs: Other args to pass through to underlying Dataset and/or Reader classes
 
     Returns:
         Dataset object
@@ -162,6 +165,7 @@ def create_dataset(
             split=split,
             class_map=class_map,
             input_img_mode=input_img_mode,
+            trust_remote_code=trust_remote_code,
             **kwargs,
         )
     elif name.startswith('hfids/'):
@@ -177,7 +181,8 @@ def create_dataset(
             repeats=repeats,
             seed=seed,
             input_img_mode=input_img_mode,
-            **kwargs
+            trust_remote_code=trust_remote_code,
+            **kwargs,
         )
     elif name.startswith('tfds/'):
         ds = IterableImageDataset(

timm/data/readers/reader_hfds.py

@@ -48,7 +48,7 @@ def __init__(
         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
+            cache_dir=self.root,  # timm doesn't expect hidden cache dir for datasets, specify a path if root set
             trust_remote_code=trust_remote_code
         )
         # leave decode for caller, plus we want easy access to original path names...

timm/data/readers/reader_hfids.py

@@ -44,6 +44,7 @@ def __init__(
             target_img_mode: str = '',
             shuffle_size: Optional[int] = None,
             num_samples: Optional[int] = None,
+            trust_remote_code: bool = False
     ):
         super().__init__()
         self.root = root
@@ -60,7 +61,11 @@ def __init__(
         self.target_key = target_key
         self.target_img_mode = target_img_mode
 
-        self.builder = datasets.load_dataset_builder(name, cache_dir=root)
+        self.builder = datasets.load_dataset_builder(
+            name,
+            cache_dir=root,
+            trust_remote_code=trust_remote_code,
+        )
         if download:
             self.builder.download_and_prepare()
 

timm/models/_builder.py

@@ -2,6 +2,7 @@
 import logging
 import os
 from copy import deepcopy
+from pathlib import Path
 from typing import Any, Callable, Dict, Optional, Tuple
 from contextlib import nullcontext
 
@@ -92,6 +93,7 @@ def load_custom_pretrained(
         model: nn.Module,
         pretrained_cfg: Optional[Dict] = None,
         load_fn: Optional[Callable] = None,
+        cache_dir: Optional[Union[str, Path]] = None,
 ):
     r"""Loads a custom (read non .pth) weight file
 
@@ -104,9 +106,10 @@ def load_custom_pretrained(
 
     Args:
         model: The instantiated model to load weights into
-        pretrained_cfg (dict): Default pretrained model cfg
+        pretrained_cfg: Default pretrained model cfg
         load_fn: An external standalone fn that loads weights into provided model, otherwise a fn named
-            'laod_pretrained' on the model will be called if it exists
+            'load_pretrained' on the model will be called if it exists
+        cache_dir: Override model checkpoint cache dir for this load
     """
     pretrained_cfg = pretrained_cfg or getattr(model, 'pretrained_cfg', None)
     if not pretrained_cfg:
@@ -124,6 +127,7 @@ def load_custom_pretrained(
             pretrained_loc,
             check_hash=_CHECK_HASH,
             progress=_DOWNLOAD_PROGRESS,
+            cache_dir=cache_dir,
         )
 
     if load_fn is not None:
@@ -141,17 +145,18 @@ def load_pretrained(
         in_chans: int = 3,
         filter_fn: Optional[Callable] = None,
         strict: bool = True,
+        cache_dir: Optional[Union[str, Path]] = None,
 ):
     """ Load pretrained checkpoint
 
     Args:
-        model (nn.Module) : PyTorch model module
-        pretrained_cfg (Optional[Dict]): configuration for pretrained weights / target dataset
-        num_classes (int): num_classes for target model
-        in_chans (int): in_chans for target model
-        filter_fn (Optional[Callable]): state_dict filter fn for load (takes state_dict, model as args)
-        strict (bool): strict load of checkpoint
-
+        model: PyTorch module
+        pretrained_cfg: Configuration for pretrained weights / target dataset
+        num_classes: Number of classes for target model. Will adapt pretrained if different.
+        in_chans: Number of input chans for target model. Will adapt pretrained if different.
+        filter_fn: state_dict filter fn for load (takes state_dict, model as args)
+        strict: Strict load of checkpoint
+        cache_dir: Override model checkpoint cache dir for this load
     """
     pretrained_cfg = pretrained_cfg or getattr(model, 'pretrained_cfg', None)
     if not pretrained_cfg:
@@ -175,6 +180,7 @@ def load_pretrained(
                 pretrained_loc,
                 progress=_DOWNLOAD_PROGRESS,
                 check_hash=_CHECK_HASH,
+                cache_dir=cache_dir,
             )
             model.load_pretrained(pretrained_loc)
             return
@@ -186,25 +192,27 @@ def load_pretrained(
                     progress=_DOWNLOAD_PROGRESS,
                     check_hash=_CHECK_HASH,
                     weights_only=True,
+                    model_dir=cache_dir,
                 )
             except TypeError:
                 state_dict = load_state_dict_from_url(
                     pretrained_loc,
                     map_location='cpu',
                     progress=_DOWNLOAD_PROGRESS,
                     check_hash=_CHECK_HASH,
+                    model_dir=cache_dir,
                 )
     elif load_from == 'hf-hub':
         _logger.info(f'Loading pretrained weights from Hugging Face hub ({pretrained_loc})')
         if isinstance(pretrained_loc, (list, tuple)):
             custom_load = pretrained_cfg.get('custom_load', False)
             if isinstance(custom_load, str) and custom_load == 'hf':
-                load_custom_from_hf(*pretrained_loc, model)
+                load_custom_from_hf(*pretrained_loc, model, cache_dir=cache_dir)
                 return
             else:
-                state_dict = load_state_dict_from_hf(*pretrained_loc)
+                state_dict = load_state_dict_from_hf(*pretrained_loc, cache_dir=cache_dir)
         else:
-            state_dict = load_state_dict_from_hf(pretrained_loc, weights_only=True)
+            state_dict = load_state_dict_from_hf(pretrained_loc, weights_only=True, cache_dir=cache_dir)
     else:
         model_name = pretrained_cfg.get('architecture', 'this model')
         raise RuntimeError(f"No pretrained weights exist for {model_name}. Use `pretrained=False` for random init.")
@@ -321,8 +329,8 @@ def _update_default_model_kwargs(pretrained_cfg, kwargs, kwargs_filter):
 
 def resolve_pretrained_cfg(
         variant: str,
-        pretrained_cfg=None,
-        pretrained_cfg_overlay=None,
+        pretrained_cfg: Optional[Union[str, Dict[str, Any]]] = None,
+        pretrained_cfg_overlay: Optional[Dict[str, Any]] = None,
 ) -> PretrainedCfg:
     model_with_tag = variant
     pretrained_tag = None
@@ -364,6 +372,7 @@ def build_model_with_cfg(
         feature_cfg: Optional[Dict] = None,
         pretrained_strict: bool = True,
         pretrained_filter_fn: Optional[Callable] = None,
+        cache_dir: Optional[Union[str, Path]] = None,
         kwargs_filter: Optional[Tuple[str]] = None,
         **kwargs,
 ):
@@ -376,16 +385,18 @@ def build_model_with_cfg(
       * pruning config / model adaptation
 
     Args:
-        model_cls: model class
-        variant: model variant name
-        pretrained: load pretrained weights
-        pretrained_cfg: model's pretrained weight/task config
-        model_cfg: model's architecture config
-        feature_cfg: feature extraction adapter config
-        pretrained_strict: load pretrained weights strictly
-        pretrained_filter_fn: filter callable for pretrained weights
-        kwargs_filter: kwargs to filter before passing to model
-        **kwargs: model args passed through to model __init__
+        model_cls: Model class
+        variant: Model variant name
+        pretrained: Load the pretrained weights
+        pretrained_cfg: Model's pretrained weight/task config
+        pretrained_cfg_overlay: Entries that will override those in pretrained_cfg
+        model_cfg: Model's architecture config
+        feature_cfg: Feature extraction adapter config
+        pretrained_strict: Load pretrained weights strictly
+        pretrained_filter_fn: Filter callable for pretrained weights
+        cache_dir: Override model cache dir for Hugging Face Hub and Torch checkpoints
+        kwargs_filter: Kwargs keys to filter (remove) before passing to model
+        **kwargs: Model args passed through to model __init__
     """
     pruned = kwargs.pop('pruned', False)
     features = False
@@ -397,8 +408,6 @@ def build_model_with_cfg(
         pretrained_cfg=pretrained_cfg,
         pretrained_cfg_overlay=pretrained_cfg_overlay
     )
-
-    # FIXME converting back to dict, PretrainedCfg use should be propagated further, but not into model
     pretrained_cfg = pretrained_cfg.to_dict()
 
     _update_default_model_kwargs(pretrained_cfg, kwargs, kwargs_filter)
@@ -437,6 +446,7 @@ def build_model_with_cfg(
             in_chans=kwargs.get('in_chans', 3),
             filter_fn=pretrained_filter_fn,
             strict=pretrained_strict,
+            cache_dir=cache_dir,
         )
 
     # Wrap the model in a feature extraction module if enabled

timm/models/_factory.py

@@ -1,4 +1,5 @@
 import os
+from pathlib import Path
 from typing import Any, Dict, Optional, Union
 from urllib.parse import urlsplit
 
@@ -40,7 +41,8 @@ def create_model(
         pretrained: bool = False,
         pretrained_cfg: Optional[Union[str, Dict[str, Any], PretrainedCfg]] = None,
         pretrained_cfg_overlay:  Optional[Dict[str, Any]] = None,
-        checkpoint_path: str = '',
+        checkpoint_path: Optional[Union[str, Path]] = None,
+        cache_dir: Optional[Union[str, Path]] = None,
         scriptable: Optional[bool] = None,
         exportable: Optional[bool] = None,
         no_jit: Optional[bool] = None,
@@ -50,17 +52,17 @@ def create_model(
 
     Lookup model's entrypoint function and pass relevant args to create a new model.
 
-    <Tip>
+    Tip:
         **kwargs will be passed through entrypoint fn to ``timm.models.build_model_with_cfg()``
         and then the model class __init__(). kwargs values set to None are pruned before passing.
-    </Tip>
 
     Args:
         model_name: Name of model to instantiate.
         pretrained: If set to `True`, load pretrained ImageNet-1k weights.
         pretrained_cfg: Pass in an external pretrained_cfg for model.
         pretrained_cfg_overlay: Replace key-values in base pretrained_cfg with these.
         checkpoint_path: Path of checkpoint to load _after_ the model is initialized.
+        cache_dir: Override model cache dir for Hugging Face Hub and Torch checkpoints.
         scriptable: Set layer config so that model is jit scriptable (not working for all models yet).
         exportable: Set layer config so that model is traceable / ONNX exportable (not fully impl/obeyed yet).
         no_jit: Set layer config so that model doesn't utilize jit scripted layers (so far activations only).
@@ -87,6 +89,10 @@ def create_model(
     >>> model = create_model('mobilenetv3_large_100', pretrained=True, num_classes=10)
     >>> model.num_classes
     10
+
+    >>> # Create a Dinov2 small model with pretrained weights and save weights in a custom directory.
+    >>> model = create_model('vit_small_patch14_dinov2.lvd142m', pretrained=True, cache_dir="/data/my-models")
+    >>> # Data will be stored at `/data/my-models/models--timm--vit_small_patch14_dinov2.lvd142m/`
     ```
     """
     # Parameters that aren't supported by all models or are intended to only override model defaults if set
@@ -99,7 +105,10 @@ def create_model(
         assert not pretrained_cfg, 'pretrained_cfg should not be set when sourcing model from Hugging Face Hub.'
         # For model names specified in the form `hf-hub:path/architecture_name@revision`,
         # load model weights + pretrained_cfg from Hugging Face hub.
-        pretrained_cfg, model_name, model_args = load_model_config_from_hf(model_name)
+        pretrained_cfg, model_name, model_args = load_model_config_from_hf(
+            model_name,
+            cache_dir=cache_dir,
+        )
         if model_args:
             for k, v in model_args.items():
                 kwargs.setdefault(k, v)
@@ -118,6 +127,7 @@ def create_model(
             pretrained=pretrained,
             pretrained_cfg=pretrained_cfg,
             pretrained_cfg_overlay=pretrained_cfg_overlay,
+            cache_dir=cache_dir,
             **kwargs,
         )
 

timm/models/_helpers.py

@@ -4,7 +4,6 @@
 """
 import logging
 import os
-from collections import OrderedDict
 from typing import Any, Callable, Dict, Optional, Union
 
 import torch