Spaces:
Build error
Build error
import numpy as np | |
import torch | |
import torch.nn as nn | |
from collections import OrderedDict | |
def tf2th(conv_weights): | |
"""Possibly convert HWIO to OIHW.""" | |
if conv_weights.ndim == 4: | |
conv_weights = conv_weights.transpose([3, 2, 0, 1]) | |
return torch.from_numpy(conv_weights) | |
def _rename_conv_weights_for_deformable_conv_layers(state_dict, cfg): | |
import re | |
layer_keys = sorted(state_dict.keys()) | |
for ix, stage_with_dcn in enumerate(cfg.MODEL.RESNETS.STAGE_WITH_DCN, 1): | |
if not stage_with_dcn: | |
continue | |
for old_key in layer_keys: | |
pattern = ".*block{}.*conv2.*".format(ix) | |
r = re.match(pattern, old_key) | |
if r is None: | |
continue | |
for param in ["weight", "bias"]: | |
if old_key.find(param) is -1: | |
continue | |
if 'unit01' in old_key: | |
continue | |
new_key = old_key.replace( | |
"conv2.{}".format(param), "conv2.conv.{}".format(param) | |
) | |
print("pattern: {}, old_key: {}, new_key: {}".format( | |
pattern, old_key, new_key | |
)) | |
# Calculate SD conv weight | |
w = state_dict[old_key] | |
v, m = torch.var_mean(w, dim=[1, 2, 3], keepdim=True, unbiased=False) | |
w = (w - m) / torch.sqrt(v + 1e-10) | |
state_dict[new_key] = w | |
del state_dict[old_key] | |
return state_dict | |
def load_big_format(cfg, f): | |
model = OrderedDict() | |
weights = np.load(f) | |
cmap = {'a':1, 'b':2, 'c':3} | |
for key, val in weights.items(): | |
old_key = key.replace('resnet/', '') | |
if 'root_block' in old_key: | |
new_key = 'root.conv.weight' | |
elif '/proj/standardized_conv2d/kernel' in old_key: | |
key_pattern = old_key.replace('/proj/standardized_conv2d/kernel', '').replace('resnet/', '') | |
bname, uname, cidx = key_pattern.split('/') | |
new_key = '{}.downsample.{}.conv{}.weight'.format(bname,uname,cmap[cidx]) | |
elif '/standardized_conv2d/kernel' in old_key: | |
key_pattern = old_key.replace('/standardized_conv2d/kernel', '').replace('resnet/', '') | |
bname, uname, cidx = key_pattern.split('/') | |
new_key = '{}.{}.conv{}.weight'.format(bname,uname,cmap[cidx]) | |
elif '/group_norm/gamma' in old_key: | |
key_pattern = old_key.replace('/group_norm/gamma', '').replace('resnet/', '') | |
bname, uname, cidx = key_pattern.split('/') | |
new_key = '{}.{}.gn{}.weight'.format(bname,uname,cmap[cidx]) | |
elif '/group_norm/beta' in old_key: | |
key_pattern = old_key.replace('/group_norm/beta', '').replace('resnet/', '') | |
bname, uname, cidx = key_pattern.split('/') | |
new_key = '{}.{}.gn{}.bias'.format(bname,uname,cmap[cidx]) | |
else: | |
print('Unknown key {}'.format(old_key)) | |
continue | |
print('Map {} -> {}'.format(key, new_key)) | |
model[new_key] = tf2th(val) | |
model = _rename_conv_weights_for_deformable_conv_layers(model, cfg) | |
return dict(model=model) | |