Add CPUBone model

tests/test_cpubone.py

@@ -0,0 +1,64 @@
+import pytest
+import torch
+
+import timm
+from timm.models.cpubone_original import CPUBoneBackbone as CPUBoneBackboneOrig
+from timm.models.cpubone import CPUBone
+
+
+def test_backbone_equal():
+    torch.manual_seed(42)
+    # stem/stages are created first and in the same order as the original backbone,
+    # so their weights match the seeded original even though a head is created afterwards
+    model = CPUBone([16, 32, 64, 128, 256], [0, 1, 1, 3, 4])
+    torch.manual_seed(42)
+    backbone_orig = CPUBoneBackboneOrig([16, 32, 64, 128, 256], [0, 1, 1, 3, 4])
+
+    torch.manual_seed(42)
+    x = torch.randn(1, 3, 224, 224)
+
+    # reconstruct the original per-stage output dict by walking stem/stages
+    torch.manual_seed(42)
+    out = {}
+    feat = model.stem(x)
+    out["stage0"] = feat
+    num_stages = len(model.stages)
+    for stage_num, stage in enumerate(model.stages, start=1):
+        feat = stage(feat)
+        key = "stage_final" if stage_num == num_stages else f"stage{stage_num}"
+        out[key] = feat
+
+    torch.manual_seed(42)
+    out_orig = backbone_orig(x)
+
+    print(out.keys(), out_orig.keys())
+    for key in out.keys():
+        out_lay = out[key]
+        out_lay_orig = out_orig[key]
+        print(key, out_lay.shape, out_lay_orig.shape)
+        # allclose instead of equal: F.scaled_dot_product_attention uses fused kernels whose
+        # float accumulation order differs from the original manual implementation (~1e-5 max diff)
+        assert torch.allclose(out_lay, out_lay_orig, atol=1e-4), f"Equality test failed in output '{key}'"
+
+
+def test_create_model():
+    model = timm.create_model("cpubone_b0")
+    model.eval()
+
+    x = torch.randn(1, 3, 224, 224)
+    with torch.no_grad():
+        out = model(x)
+
+    assert out.shape == (1, 1000)
+    assert not torch.isnan(out).any()
+
+
+def test_create_model_num_classes():
+    model = timm.create_model("cpubone_nano", num_classes=10)
+    model.eval()
+
+    x = torch.randn(1, 3, 224, 224)
+    with torch.no_grad():
+        out = model(x)
+
+    assert out.shape == (1, 10)

tests/test_models.py

@@ -57,7 +57,7 @@
     'tiny_vit', 'vovnet', 'tresnet', 'rexnet', 'resnetv2', 'repghost', 'repvit', 'pvt_v2', 'nextvit', 'nest',
     'mambaout', 'inception_next', 'inception_v4', 'hgnet', 'gcvit', 'focalnet', 'efficientformer_v2', 'edgenext',
     'davit', 'rdnet', 'convnext', 'pit', 'starnet', 'shvit', 'fasternet', 'swiftformer', 'ghostnet', 'naflexvit',
-    'csatv2'
+    'csatv2', 'cpubone'
 ]
 
 # transformer / hybrid models don't support full set of spatial / feature APIs and/or have spatial output.
@@ -244,6 +244,7 @@ def test_model_backward(model_name, batch_size):
 
 # models with extra conv/linear layers after pooling
 EARLY_POOL_MODELS = (
+    timm.models.CPUBone,
     timm.models.EfficientVit,
     timm.models.EfficientVitLarge,
     timm.models.FasterNet,

timm/models/__init__.py

@@ -6,6 +6,7 @@
 from .convit import *
 from .convmixer import *
 from .convnext import *
+from .cpubone import *
 from .crossvit import *
 from .csatv2 import *
 from .cspnet import *

timm/models/_hub.py

@@ -220,6 +220,17 @@ def load_state_dict_from_hf(
     assert has_hf_hub(True)
     hf_model_id, hf_revision = hf_split(model_id)
 
+    # Load directly via safetensors if that's what the filename specifies
+    if filename.endswith(".safetensors"):
+        assert _has_safetensors, "`pip install safetensors` to use .safetensors"
+        cached_safe_file = hf_hub_download(
+            repo_id=hf_model_id,
+            filename=filename,
+            revision=hf_revision,
+            cache_dir=cache_dir,
+        )
+        return safetensors.torch.load_file(cached_safe_file, device="cpu")
+
     # Look for .safetensors alternatives and load from it if it exists
     if _has_safetensors:
         for safe_filename in _get_safe_alternatives(filename):