updated per new model code

pt_resnet_to_flax.py

@@ -1,23 +1,23 @@
-from transformers import ResNetConfig, FlaxResNetForImageClassification, ResNetForImageClassification
+from transformers import ResNetConfig, FlaxResNetForImageClassification, ResNetForImageClassification, FlaxResNetModel, ResNetModel
 from flax.traverse_util import flatten_dict, unflatten_dict
 from flax.core.frozen_dict import unfreeze
 import re
+import jax.numpy as jnp
+import torch
 
-pt_resnet = ResNetForImageClassification.from_pretrained("microsoft/resnet-50")
-pt_state = pt_resnet.state_dict()
 
+pt_resnet = ResNetForImageClassification.from_pretrained("microsoft/resnet-50")
 flax_resnet = FlaxResNetForImageClassification(pt_resnet.config)
+
+pt_state = pt_resnet.state_dict()
 flax_state = flatten_dict(unfreeze(flax_resnet.params))
-    
+
+
 new_pt_state = {}
-pt_batch_stats = {}
 for key, tensor in pt_state.items():
     key_parts = set(key.split("."))
     tensor = tensor.numpy()
 
-    key = re.sub(r"(?<=[a-zA-Z]).(?=\d)", "_", key)
-    
-
     if "convolution.weight" in key:
         key = key.replace("weight", "kernel")
         tensor = tensor.transpose((2, 3, 1, 0))
@@ -34,36 +34,32 @@ for key, tensor in pt_state.items():
         key = "params."+key
         new_pt_state[key] = tensor
     
-    elif "classifier_1.weight" in key:
-        key = "params.classifier.kernel"
+    elif "classifier.1.weight" in key:
+        key = "params.classifier.1.kernel"
         new_pt_state[key] = tensor.transpose()
 
-    elif "classifier_1.bias" in key:
-        key = "params.classifier.bias"
+    elif "classifier.1.bias" in key:
+        key = "params.classifier.1.bias"
         new_pt_state[key] = tensor
 
     elif "normalization.running_mean" in key:
         key = key.replace("running_mean", "mean")
-        pt_batch_stats[key] = tensor
+        key = "batch_stats."+key
+        new_pt_state[key] = tensor
 
     elif "normalization.running_var" in key:
         key = key.replace("running_var", "var")
-        pt_batch_stats[key] = tensor
+        key = "batch_stats."+key
+        new_pt_state[key] = tensor
     
     else:
-        pass
+        continue
 
+    
 for total_updated, (new_key, new_tensor) in enumerate(new_pt_state.items()):
     orig_flax_tensor = flax_state.get(tuple(new_key.split(".")))
     assert orig_flax_tensor is not None
-     
-    if not("classifier" in new_key):
-        assert orig_flax_tensor.shape == new_tensor.shape
+    assert orig_flax_tensor.shape == new_tensor.shape
     flax_state[tuple(new_key.split("."))] = new_tensor
-
 flax_state = unflatten_dict(flax_state)
-
-pt_batch_stats = unflatten_dict({tuple(k.split(".")):v for k,v in pt_batch_stats.items()})
-flax_state["batch_stats"] = pt_batch_stats
-
-flax_resnet.save_pretrained("resnet_50_flax", params=flax_state)
+flax_resnet.save_pretrained("resnet_50_flax", params=flax_state)