| 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") | |
| flax_resnet = FlaxResNetForImageClassification(pt_resnet.config) | |
| pt_state = pt_resnet.state_dict() | |
| flax_state = flatten_dict(unfreeze(flax_resnet.params)) | |
| new_pt_state = {} | |
| for key, tensor in pt_state.items(): | |
| key_parts = set(key.split(".")) | |
| tensor = tensor.numpy() | |
| if "convolution.weight" in key: | |
| key = key.replace("weight", "kernel") | |
| tensor = tensor.transpose((2, 3, 1, 0)) | |
| key = "params."+key | |
| new_pt_state[key] = tensor | |
| elif "normalization.weight" in key: | |
| key = key.replace("weight", "scale") | |
| key = "params."+key | |
| new_pt_state[key] = tensor | |
| elif "normalization.bias" in key: | |
| key = key.replace("bias", "bias") | |
| key = "params."+key | |
| new_pt_state[key] = tensor | |
| 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.1.bias" | |
| new_pt_state[key] = tensor | |
| elif "normalization.running_mean" in key: | |
| key = key.replace("running_mean", "mean") | |
| key = "batch_stats."+key | |
| new_pt_state[key] = tensor | |
| elif "normalization.running_var" in key: | |
| key = key.replace("running_var", "var") | |
| key = "batch_stats."+key | |
| new_pt_state[key] = tensor | |
| else: | |
| 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 | |
| assert orig_flax_tensor.shape == new_tensor.shape | |
| flax_state[tuple(new_key.split("."))] = new_tensor | |
| flax_state = unflatten_dict(flax_state) | |
| flax_resnet.save_pretrained("resnet_50_flax", params=flax_state) | |