refactor weights loading without using pickled files - vgg, resnet (#21)

ivy_models/__init__.py

@@ -2,7 +2,9 @@
 from . import resnet
 from . import vgg
 from . import convnext
-from .transformers import *
+from . import helpers
+from .transformers import perceiver_io
+from .transformers.perceiver_io import *
 from .resnet import *
 from .vgg import *
 from .convnext import *

ivy_models/convnext/convnext.py

@@ -1,4 +1,5 @@
 import ivy
+import ivy_models
 from ivy.stateful.module import Module
 from ivy.stateful.initializers import Zeros, Ones, Constant
 
@@ -158,6 +159,25 @@ def _forward(self, x):
         return self.v.weight[:, None, None] * x + self.v.bias[:, None, None]
 
 
+def _convnext_torch_weights_mapping(old_key, new_key):
+    new_mapping = new_key
+    if "downsample_layers" in old_key:
+        if "0/0/bias" in old_key:
+            new_mapping = {"key_chain": new_key, "pattern": "h -> 1 h 1 1"}
+        elif "0/0/weight" in old_key:
+            new_mapping = {"key_chain": new_key, "pattern": "b c h w-> h w c b"}
+        elif "downsample_layers/0" not in old_key and "1/bias" in old_key:
+            new_mapping = {"key_chain": new_key, "pattern": "h -> 1 h 1 1"}
+        elif "downsample_layers/0" not in old_key and "1/w" in old_key:
+            new_mapping = {"key_chain": new_key, "pattern": "b c h w -> h w c b"}
+    elif "dwconv" in old_key:
+        if "bias" in old_key:
+            new_mapping = {"key_chain": new_key, "pattern": "h -> 1 h 1 1"}
+        elif "weight" in old_key:
+            new_mapping = {"key_chain": new_key, "pattern": "a 1 c d -> c d a"}
+    return new_mapping
+
+
 def convnext(size: str, pretrained=True):
     """Loads a ConvNeXt with specified size, optionally pretrained."""
     size_dict = {
@@ -181,36 +201,8 @@ def convnext(size: str, pretrained=True):
         "large": "https://dl.fbaipublicfiles.com/convnext/convnext_large_1k_224_ema.pth",  # noqa
     }
 
-    import torch
-
-    old_backend = ivy.backend
-    ivy.set_backend("torch")
-    weights = torch.hub.load_state_dict_from_url(weight_dl[size])
-    weights_raw = ivy.to_numpy(ivy.Container(weights))
     reference_model = ConvNeXt(depths=depths, dims=dims)
-    mapping = {}
-    for old_key, new_key in zip(
-        weights_raw.cont_sort_by_key().cont_to_iterator_keys(),
-        reference_model.v.cont_sort_by_key().cont_to_iterator_keys(),
-    ):
-        new_mapping = new_key
-        if "downsample_layers" in old_key:
-            if "0/0/bias" in old_key:
-                new_mapping = {"key_chain": new_key, "pattern": "h -> 1 h 1 1"}
-            elif "0/0/weight" in old_key:
-                new_mapping = {"key_chain": new_key, "pattern": "b c h w-> h w c b"}
-            elif "downsample_layers/0" not in old_key and "1/bias" in old_key:
-                new_mapping = {"key_chain": new_key, "pattern": "h -> 1 h 1 1"}
-            elif "downsample_layers/0" not in old_key and "1/w" in old_key:
-                new_mapping = {"key_chain": new_key, "pattern": "b c h w -> h w c b"}
-        elif "dwconv" in old_key:
-            if "bias" in old_key:
-                new_mapping = {"key_chain": new_key, "pattern": "h -> 1 h 1 1"}
-            elif "weight" in old_key:
-                new_mapping = {"key_chain": new_key, "pattern": "a 1 c d -> c d a"}
-        mapping[old_key] = new_mapping
-
-    ivy.set_backend(old_backend)
-    w_clean = weights_raw.cont_restructure(mapping, keep_orig=False)
-    w_clean = ivy.asarray(w_clean)
+    w_clean = ivy_models.helpers.load_torch_weights(
+        weight_dl[size], reference_model, custom_mapping=_convnext_torch_weights_mapping
+    )
     return ConvNeXt(depths=depths, dims=dims, v=w_clean)

ivy_models/helpers/__init__.py

@@ -1 +1 @@
-from .match_weights_helpers import *
+from .weights_helpers import *

ivy_models/helpers/weights_helpers.py

@@ -0,0 +1,69 @@
+# global
+import ivy
+
+
+def _prune_keys(raw, ref, raw_keys_to_prune=[], ref_keys_to_prune=[]):
+    if raw_keys_to_prune != []:
+        for kc in raw_keys_to_prune:
+            raw = raw.cont_prune_key_from_key_chains(containing=kc)
+    if ref_keys_to_prune != []:
+        for kc in ref_keys_to_prune:
+            ref = ref.cont_prune_key_from_key_chains(containing=kc)
+    return raw, ref
+
+
+def _map_weights(raw, ref, custom_mapping=None):
+    mapping = {}
+    for old_key, new_key in zip(
+        raw.cont_sort_by_key().cont_to_iterator_keys(),
+        ref.cont_sort_by_key().cont_to_iterator_keys(),
+    ):
+        new_mapping = new_key
+        if custom_mapping is not None:
+            new_mapping = custom_mapping(old_key, new_key)
+            if new_mapping is None:
+                continue
+        mapping[old_key] = new_mapping
+    return mapping
+
+
+def load_torch_weights(url, ref_model, custom_mapping=None):
+    import torch
+
+    ivy.set_backend("torch")
+    weights = torch.hub.load_state_dict_from_url(url)
+    weights_raw = ivy.to_numpy(ivy.Container(weights))
+    mapping = _map_weights(weights_raw, ref_model.v, custom_mapping=custom_mapping)
+
+    ivy.previous_backend()
+    w_clean = weights_raw.cont_restructure(mapping, keep_orig=False)
+    return ivy.asarray(w_clean)
+
+
+def load_jax_weights(
+    url, ref_model, custom_mapping=None, raw_keys_to_prune=[], ref_keys_to_prune=[]
+):
+    import urllib.request
+    import os
+    import pickle
+
+    ivy.set_backend("jax")
+    urllib.request.urlretrieve(url, filename="jax_weights.pystate")
+    with open("jax_weights.pystate", "rb") as f:
+        weights = pickle.loads(f.read())
+    os.remove("jax_weights.pystate")
+
+    try:
+        weights = {**weights["params"], **weights["state"]}
+    except KeyError:
+        pass
+
+    weights_raw = ivy.to_numpy(ivy.Container(weights))
+    weights_raw, weights_ref = _prune_keys(
+        weights_raw, ref_model.v, raw_keys_to_prune, ref_keys_to_prune
+    )
+    mapping = _map_weights(weights_raw, weights_ref, custom_mapping=custom_mapping)
+
+    ivy.previous_backend()
+    w_clean = weights_raw.cont_restructure(mapping, keep_orig=False)
+    return ivy.asarray(w_clean)

ivy_models/resnet/resnet.py

@@ -1,5 +1,7 @@
 # global
+import builtins
 import ivy
+import ivy_models
 
 
 class ResidualBlock(ivy.Module):
@@ -149,9 +151,27 @@ def _forward(self, x):
         return x
 
 
-def resnet_18(v=None):
+def _resnet_torch_weights_mapping(old_key, new_key):
+    W_KEY = ["conv1/weight", "conv2/weight", "downsample/0/weight"]
+    new_mapping = new_key
+    if builtins.any([kc in old_key for kc in W_KEY]):
+        new_mapping = {"key_chain": new_key, "pattern": "b c h w -> h w c b"}
+    elif "num_batches_tracked" in old_key:
+        new_mapping = None
+    return new_mapping
+
+
+def resnet_18(pretrained=True):
     """ResNet-18 model"""
-    return ResNet(ResidualBlock, [2, 2, 2, 2], v=v)
+    if not pretrained:
+        return ResNet(ResidualBlock, [2, 2, 2, 2])
+
+    reference_model = ResNet(ResidualBlock, [2, 2, 2, 2])
+    url = "https://download.pytorch.org/models/resnet18-f37072fd.pth"
+    w_clean = ivy_models.helpers.load_torch_weights(
+        url, reference_model, custom_mapping=_resnet_torch_weights_mapping
+    )
+    return ResNet(ResidualBlock, [2, 2, 2, 2], v=w_clean)
 
 
 def resnet_34(v=None):