Add a deprecation phase to module re-org

Author: rwightman

benchmark.py

@@ -19,7 +19,8 @@
 import torch.nn.parallel
 
 from timm.data import resolve_data_config
-from timm.models import create_model, is_model, list_models, set_fast_norm
+from timm.layers import set_fast_norm
+from timm.models import create_model, is_model, list_models
 from timm.optim import create_optimizer_v2
 from timm.utils import setup_default_logging, set_jit_fuser, decay_batch_step, check_batch_size_retry
 

timm/models/_builder.py

@@ -23,6 +23,10 @@
 _CHECK_HASH = False
 
 
+__all__ = ['set_pretrained_download_progress', 'set_pretrained_check_hash', 'load_custom_pretrained', 'load_pretrained',
+           'pretrained_cfg_for_features', 'resolve_pretrained_cfg', 'build_model_with_cfg']
+
+
 def _resolve_pretrained_source(pretrained_cfg):
     cfg_source = pretrained_cfg.get('source', '')
     pretrained_url = pretrained_cfg.get('url', None)

timm/models/_factory.py

@@ -9,6 +9,9 @@
 from ._registry import is_model, model_entrypoint
 
 
+__all__ = ['parse_model_name', 'safe_model_name', 'create_model']
+
+
 def parse_model_name(model_name):
     if model_name.startswith('hf_hub'):
         # NOTE for backwards compat, deprecate hf_hub use

timm/models/_features.py

@@ -17,6 +17,9 @@
 import torch.nn as nn
 
 
+__all__ = ['FeatureInfo', 'FeatureHooks', 'FeatureDictNet', 'FeatureListNet', 'FeatureHookNet']
+
+
 class FeatureInfo:
 
     def __init__(self, feature_info: List[Dict], out_indices: Tuple[int]):

timm/models/_features_fx.py

@@ -35,6 +35,10 @@
     pass
 
 
+__all__ = ['register_notrace_module', 'register_notrace_function', 'create_feature_extractor',
+           'FeatureGraphNet', 'GraphExtractNet']
+
+
 def register_notrace_module(module: Type[nn.Module]):
     """
     Any module not under timm.models.layers should get this decorator if we don't want to trace through it.

timm/models/_helpers.py

@@ -12,6 +12,8 @@
 
 _logger = logging.getLogger(__name__)
 
+__all__ = ['clean_state_dict', 'load_state_dict', 'load_checkpoint', 'remap_checkpoint', 'resume_checkpoint']
+
 
 def clean_state_dict(state_dict):
     # 'clean' checkpoint by removing .module prefix from state dict if it exists from parallel training

timm/models/_hub.py

@@ -31,6 +31,9 @@
 
 _logger = logging.getLogger(__name__)
 
+__all__ = ['get_cache_dir', 'download_cached_file', 'has_hf_hub', 'hf_split', 'load_model_config_from_hf',
+           'load_state_dict_from_hf', 'save_for_hf', 'push_to_hf_hub']
+
 
 def get_cache_dir(child_dir=''):
     """

timm/models/_manipulate.py

@@ -9,6 +9,9 @@
 from torch import nn as nn
 from torch.utils.checkpoint import checkpoint
 
+__all__ = ['model_parameters', 'named_apply', 'named_modules', 'named_modules_with_params', 'adapt_input_conv',
+           'group_with_matcher', 'group_modules', 'group_parameters', 'flatten_modules', 'checkpoint_seq']
+
 
 def model_parameters(model, exclude_head=False):
     if exclude_head:

timm/models/_pretrained.py

@@ -4,6 +4,9 @@
 from typing import Any, Deque, Dict, Tuple, Optional, Union
 
 
+__all__ = ['PretrainedCfg', 'filter_pretrained_cfg', 'DefaultCfg', 'split_model_name_tag', 'generate_default_cfgs']
+
+
 @dataclass
 class PretrainedCfg:
     """

timm/models/_prune.py

@@ -5,6 +5,8 @@
 
 from timm.layers import Conv2dSame, BatchNormAct2d, Linear
 
+__all__ = ['extract_layer', 'set_layer', 'adapt_model_from_string', 'adapt_model_from_file']
+
 
 def extract_layer(model, layer):
     layer = layer.split('.')