Spaces:
Sleeping
Sleeping
Segment-Anything
/
transformers_4_35_0
/models
/bert
/convert_bert_original_tf2_checkpoint_to_pytorch.py
# Copyright 2020 The HuggingFace Team. All rights reserved. | |
# | |
# Licensed under the Apache License, Version 2.0 (the "License"); | |
# you may not use this file except in compliance with the License. | |
# You may obtain a copy of the License at | |
# | |
# http://www.apache.org/licenses/LICENSE-2.0 | |
# | |
# Unless required by applicable law or agreed to in writing, software | |
# distributed under the License is distributed on an "AS IS" BASIS, | |
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | |
# See the License for the specific language governing permissions and | |
# limitations under the License. | |
""" | |
This script can be used to convert a head-less TF2.x Bert model to PyTorch, as published on the official (now | |
deprecated) GitHub: https://github.com/tensorflow/models/tree/v2.3.0/official/nlp/bert | |
TF2.x uses different variable names from the original BERT (TF 1.4) implementation. The script re-maps the TF2.x Bert | |
weight names to the original names, so the model can be imported with Huggingface/transformer. | |
You may adapt this script to include classification/MLM/NSP/etc. heads. | |
Note: This script is only working with an older version of the TensorFlow models repository (<= v2.3.0). | |
Models trained with never versions are not compatible with this script. | |
""" | |
import argparse | |
import os | |
import re | |
import tensorflow as tf | |
import torch | |
from transformers import BertConfig, BertModel | |
from transformers.utils import logging | |
logging.set_verbosity_info() | |
logger = logging.get_logger(__name__) | |
def load_tf2_weights_in_bert(model, tf_checkpoint_path, config): | |
tf_path = os.path.abspath(tf_checkpoint_path) | |
logger.info(f"Converting TensorFlow checkpoint from {tf_path}") | |
# Load weights from TF model | |
init_vars = tf.train.list_variables(tf_path) | |
names = [] | |
arrays = [] | |
layer_depth = [] | |
for full_name, shape in init_vars: | |
# logger.info(f"Loading TF weight {name} with shape {shape}") | |
name = full_name.split("/") | |
if full_name == "_CHECKPOINTABLE_OBJECT_GRAPH" or name[0] in ["global_step", "save_counter"]: | |
logger.info(f"Skipping non-model layer {full_name}") | |
continue | |
if "optimizer" in full_name: | |
logger.info(f"Skipping optimization layer {full_name}") | |
continue | |
if name[0] == "model": | |
# ignore initial 'model' | |
name = name[1:] | |
# figure out how many levels deep the name is | |
depth = 0 | |
for _name in name: | |
if _name.startswith("layer_with_weights"): | |
depth += 1 | |
else: | |
break | |
layer_depth.append(depth) | |
# read data | |
array = tf.train.load_variable(tf_path, full_name) | |
names.append("/".join(name)) | |
arrays.append(array) | |
logger.info(f"Read a total of {len(arrays):,} layers") | |
# Sanity check | |
if len(set(layer_depth)) != 1: | |
raise ValueError(f"Found layer names with different depths (layer depth {list(set(layer_depth))})") | |
layer_depth = list(set(layer_depth))[0] | |
if layer_depth != 1: | |
raise ValueError( | |
"The model contains more than just the embedding/encoder layers. This script does not handle MLM/NSP" | |
" heads." | |
) | |
# convert layers | |
logger.info("Converting weights...") | |
for full_name, array in zip(names, arrays): | |
name = full_name.split("/") | |
pointer = model | |
trace = [] | |
for i, m_name in enumerate(name): | |
if m_name == ".ATTRIBUTES": | |
# variable names end with .ATTRIBUTES/VARIABLE_VALUE | |
break | |
if m_name.startswith("layer_with_weights"): | |
layer_num = int(m_name.split("-")[-1]) | |
if layer_num <= 2: | |
# embedding layers | |
# layer_num 0: word_embeddings | |
# layer_num 1: position_embeddings | |
# layer_num 2: token_type_embeddings | |
continue | |
elif layer_num == 3: | |
# embedding LayerNorm | |
trace.extend(["embeddings", "LayerNorm"]) | |
pointer = getattr(pointer, "embeddings") | |
pointer = getattr(pointer, "LayerNorm") | |
elif layer_num > 3 and layer_num < config.num_hidden_layers + 4: | |
# encoder layers | |
trace.extend(["encoder", "layer", str(layer_num - 4)]) | |
pointer = getattr(pointer, "encoder") | |
pointer = getattr(pointer, "layer") | |
pointer = pointer[layer_num - 4] | |
elif layer_num == config.num_hidden_layers + 4: | |
# pooler layer | |
trace.extend(["pooler", "dense"]) | |
pointer = getattr(pointer, "pooler") | |
pointer = getattr(pointer, "dense") | |
elif m_name == "embeddings": | |
trace.append("embeddings") | |
pointer = getattr(pointer, "embeddings") | |
if layer_num == 0: | |
trace.append("word_embeddings") | |
pointer = getattr(pointer, "word_embeddings") | |
elif layer_num == 1: | |
trace.append("position_embeddings") | |
pointer = getattr(pointer, "position_embeddings") | |
elif layer_num == 2: | |
trace.append("token_type_embeddings") | |
pointer = getattr(pointer, "token_type_embeddings") | |
else: | |
raise ValueError(f"Unknown embedding layer with name {full_name}") | |
trace.append("weight") | |
pointer = getattr(pointer, "weight") | |
elif m_name == "_attention_layer": | |
# self-attention layer | |
trace.extend(["attention", "self"]) | |
pointer = getattr(pointer, "attention") | |
pointer = getattr(pointer, "self") | |
elif m_name == "_attention_layer_norm": | |
# output attention norm | |
trace.extend(["attention", "output", "LayerNorm"]) | |
pointer = getattr(pointer, "attention") | |
pointer = getattr(pointer, "output") | |
pointer = getattr(pointer, "LayerNorm") | |
elif m_name == "_attention_output_dense": | |
# output attention dense | |
trace.extend(["attention", "output", "dense"]) | |
pointer = getattr(pointer, "attention") | |
pointer = getattr(pointer, "output") | |
pointer = getattr(pointer, "dense") | |
elif m_name == "_output_dense": | |
# output dense | |
trace.extend(["output", "dense"]) | |
pointer = getattr(pointer, "output") | |
pointer = getattr(pointer, "dense") | |
elif m_name == "_output_layer_norm": | |
# output dense | |
trace.extend(["output", "LayerNorm"]) | |
pointer = getattr(pointer, "output") | |
pointer = getattr(pointer, "LayerNorm") | |
elif m_name == "_key_dense": | |
# attention key | |
trace.append("key") | |
pointer = getattr(pointer, "key") | |
elif m_name == "_query_dense": | |
# attention query | |
trace.append("query") | |
pointer = getattr(pointer, "query") | |
elif m_name == "_value_dense": | |
# attention value | |
trace.append("value") | |
pointer = getattr(pointer, "value") | |
elif m_name == "_intermediate_dense": | |
# attention intermediate dense | |
trace.extend(["intermediate", "dense"]) | |
pointer = getattr(pointer, "intermediate") | |
pointer = getattr(pointer, "dense") | |
elif m_name == "_output_layer_norm": | |
# output layer norm | |
trace.append("output") | |
pointer = getattr(pointer, "output") | |
# weights & biases | |
elif m_name in ["bias", "beta"]: | |
trace.append("bias") | |
pointer = getattr(pointer, "bias") | |
elif m_name in ["kernel", "gamma"]: | |
trace.append("weight") | |
pointer = getattr(pointer, "weight") | |
else: | |
logger.warning(f"Ignored {m_name}") | |
# for certain layers reshape is necessary | |
trace = ".".join(trace) | |
if re.match(r"(\S+)\.attention\.self\.(key|value|query)\.(bias|weight)", trace) or re.match( | |
r"(\S+)\.attention\.output\.dense\.weight", trace | |
): | |
array = array.reshape(pointer.data.shape) | |
if "kernel" in full_name: | |
array = array.transpose() | |
if pointer.shape == array.shape: | |
pointer.data = torch.from_numpy(array) | |
else: | |
raise ValueError( | |
f"Shape mismatch in layer {full_name}: Model expects shape {pointer.shape} but layer contains shape:" | |
f" {array.shape}" | |
) | |
logger.info(f"Successfully set variable {full_name} to PyTorch layer {trace}") | |
return model | |
def convert_tf2_checkpoint_to_pytorch(tf_checkpoint_path, config_path, pytorch_dump_path): | |
# Instantiate model | |
logger.info(f"Loading model based on config from {config_path}...") | |
config = BertConfig.from_json_file(config_path) | |
model = BertModel(config) | |
# Load weights from checkpoint | |
logger.info(f"Loading weights from checkpoint {tf_checkpoint_path}...") | |
load_tf2_weights_in_bert(model, tf_checkpoint_path, config) | |
# Save pytorch-model | |
logger.info(f"Saving PyTorch model to {pytorch_dump_path}...") | |
torch.save(model.state_dict(), pytorch_dump_path) | |
if __name__ == "__main__": | |
parser = argparse.ArgumentParser() | |
parser.add_argument( | |
"--tf_checkpoint_path", type=str, required=True, help="Path to the TensorFlow 2.x checkpoint path." | |
) | |
parser.add_argument( | |
"--bert_config_file", | |
type=str, | |
required=True, | |
help="The config json file corresponding to the BERT model. This specifies the model architecture.", | |
) | |
parser.add_argument( | |
"--pytorch_dump_path", | |
type=str, | |
required=True, | |
help="Path to the output PyTorch model (must include filename).", | |
) | |
args = parser.parse_args() | |
convert_tf2_checkpoint_to_pytorch(args.tf_checkpoint_path, args.bert_config_file, args.pytorch_dump_path) | |