import numpy as np
|
import logging
|
from typing import Dict
|
import sys
|
import torch
|
|
|
def load_ckpt(checkpoint_path: str) -> Dict[str, np.ndarray]:
|
from tensorflow.python import pywrap_tensorflow
|
|
reader = pywrap_tensorflow.NewCheckpointReader(checkpoint_path)
|
var_to_shape_map = reader.get_variable_to_shape_map()
|
|
var_dict = dict()
|
for var_name in sorted(var_to_shape_map):
|
if "optimizer" in var_name:
|
continue
|
tensor = reader.get_tensor(var_name)
|
print("in ckpt: {}, {}".format(var_name, tensor.shape)) # print(tensor)
|
var_dict[var_name] = tensor
|
|
return var_dict
|
|
|
def convert_parameter_name_for_asv_resnet34(
|
var_dict: Dict[str, np.ndarray],
|
old_prefix: str = "EAND/speech_encoder",
|
new_prefix: str = "encoder",
|
train_steps: int = 0
|
) -> Dict[str, np.ndarray]:
|
new_dict = dict()
|
model_size = 0
|
for name, tensor in var_dict.items():
|
if not name.startswith(old_prefix):
|
if name == "softmax/output/kernel":
|
new_name = "decoder.output_dense.weight"
|
tensor = np.transpose(tensor, [1, 0])
|
new_dict[new_name] = torch.Tensor(tensor)
|
continue
|
new_name = name.replace(old_prefix, new_prefix)
|
new_name = new_name.replace("/", ".")
|
if "resnet1" in new_name or "resnet2" in new_name:
|
new_name = new_name.replace("encoder", "decoder")
|
module_name, para_name = new_name.rsplit(".", 1)
|
# process for batch normalization
|
if "bn" in module_name:
|
new_name = new_name.replace("gamma", "weight")
|
new_name = new_name.replace("beta", "bias")
|
new_name = new_name.replace("moving_mean", "running_mean")
|
new_name = new_name.replace("moving_variance", "running_var")
|
|
new_dict[new_name] = torch.Tensor(tensor)
|
new_dict[module_name + ".num_batches_tracked"] = torch.Tensor(train_steps)
|
|
# process for dense layers
|
elif "dense" in module_name:
|
new_name = new_name.replace("kernel", "weight")
|
if para_name == "kernel":
|
if len(tensor.shape) == 2:
|
tensor = np.transpose(tensor, [1, 0])
|
elif len(tensor.shape) == 3:
|
tensor = np.transpose(tensor, [2, 1, 0])
|
# for dense0
|
elif len(tensor.shape) == 4:
|
tensor = np.transpose(tensor, [3, 2, 0, 1])
|
|
new_dict[new_name] = torch.Tensor(tensor)
|
|
# process for conv layers
|
elif "conv" in module_name:
|
new_name = new_name.replace("kernel", "weight")
|
if para_name == "kernel":
|
tensor = np.transpose(tensor, [3, 2, 0, 1])
|
|
new_dict[new_name] = torch.Tensor(tensor)
|
|
print("{} -> {}".format(name, new_name))
|
model_size += new_dict[new_name].numel()
|
print("Model size: {}".format(model_size))
|
return new_dict
|
|
|
if __name__ == '__main__':
|
checkpoint_path = sys.argv[1]
|
pkl_path = sys.argv[2]
|
tf_dict = load_ckpt(checkpoint_path)
|
torch_dict = convert_parameter_name_for_asv_resnet34(
|
tf_dict,
|
train_steps=300000,
|
)
|
torch.save(
|
torch_dict, pkl_path
|
)
|
