def pytorch_to_nnabla(input_file, h5_file):
read = torch.load(input_file)
for k, v in read.items():
key = rename_params(k)
params = PF.get_parameter_or_create(key, shape=v.shape)
params.d = v.numpy()
nn.parameter.save_parameters(h5_file)
def parse_tf_ckpt(ckpt_file, h5_file):
''' Parse the TF checkpoint file and save as nnabla parameters
'''
# Get tensorflow checkpoint reader
reader = pywrap_tensorflow.NewCheckpointReader(ckpt_file)
var_to_shape_map = reader.get_variable_to_shape_map()
# Loop through each tensor name from the variable to shape map
for key in sorted(var_to_shape_map):
# Read tensor values for each tensor name
weight = reader.get_tensor(key)
if 'depthwise' in key and weight.ndim == 4:
weight = numpy.squeeze(weight, axis=3)
weight = numpy.transpose(weight, (2, 0, 1))
else:
if (weight.ndim == 4):
weight = numpy.transpose(weight, (3, 2, 0, 1))
if 'BatchNorm' in key:
weight = weight.reshape((1, -1, 1, 1))
if 'Momentum' in key or 'ExponentialMovingAverage' in key:
continue
key = rename_params(key)
# Create parameter with the same tensor name and shape
params = pf.get_parameter_or_create(key, shape=weight.shape)
params.d = weight
# Save to a h5 file
nn.parameter.save_parameters(h5_file)
def convert_ckpt_to_h5(input_file, h5_file):
"""
Convert the input checkpoint file to output hdf5 file
"""
# Get tensorflow checkpoint reader
reader = pywrap_tensorflow.NewCheckpointReader(input_file)
var_to_shape_map = reader.get_variable_to_shape_map()
# Loop through each tensor name from the variable to shape map
for key in var_to_shape_map:
# Read tensor values for each tensor name
weight = reader.get_tensor(key)
if not str(key).startswith("vgg_19/mean_rgb") and not str(
key).startswith("global_step"):
s = key.split('/')
if str(s[-2]).startswith("fc"):
k = ('/'.join(s))
key = rename_params(str(k), affine=True)
else:
s.remove(s[1])
k = ('/'.join(s))
key = rename_params(str(k), affine=False)
if weight.ndim == 4:
# transpose TF weight to NNabla weight format
weight = np.transpose(weight, (3, 0, 1, 2))
# Create parameter with the same tensor name and shape
params = PF.get_parameter_or_create(key, shape=weight.shape)
params.d = weight
# Save to a h5 file
nn.parameter.save_parameters(h5_file)
def pytorch_to_nnabla(input_file, h5_file):
read = torch.load(input_file)
for k, v in read.items():
if not str(k).startswith("classifier"):
key = rename_params(str(k), affine=False)
key = key.replace("/", "", 1)
else:
key = rename_params(str(k), affine=True)
params = pf.get_parameter_or_create(key, shape=v.shape)
params.d = v.numpy()
nn.parameter.save_parameters(h5_file)
def pytorch_to_nnabla(input_file, h5_file):
read = torch.load(input_file)
for k, v in read.items():
split = k.split('.')[-2]
if split.startswith('bn') or split.startswith('0'):
key = rename_params(k, conv=False)
v = v.reshape((1, ) + v.shape + (1, 1))
else:
key = rename_params(k, conv=True)
params = PF.get_parameter_or_create(key, shape=v.shape)
params.d = v.cpu().numpy()
nn.parameter.save_parameters(h5_file)
def pytorch_to_nnabla(input_file, h5_file):
read = torch.load(input_file)
for k, v in read['state_dict'].items():
k = k.replace('module.', '')
split = k.split('.')[-2]
if split.startswith('bn') or split.startswith('1'):
key = rename_params(k, conv=False)
v = v.reshape((1, ) + v.shape + (1, 1))
else:
if k == 'fc.weight':
v = v.T
key = rename_params(k, conv=True)
params = PF.get_parameter_or_create(key, shape=v.shape)
params.d = v.cpu().numpy()
nn.parameter.save_parameters(h5_file)
def convert(ckpt_file, h5_file):
"""
Convert BERT Tensorflow weights to NNabla
Args:
ckpt_file: Input Tensorflow ckpt file
h5_file: Output NNabla output file
"""
# Check the TensorFlow version for compatibility
if int(tf.__version__[0]) == 2:
reader = tf.compat.v1.train.NewCheckpointReader(ckpt_file)
else:
reader = pywrap_tensorflow.NewCheckpointReader(ckpt_file)
var_to_shape_map = reader.get_variable_to_shape_map()
for key in sorted(var_to_shape_map):
weight = reader.get_tensor(key)
if 'encoder' in key:
layer_id = int(key.split('/')[2].replace('layer_', ''))
if 'query/bias' in key:
key = 'encoder{:02d}/transformer_encode/src_self_attn/multi_head_attention/q_bias'.format(
layer_id)
if 'query/kernel' in key:
key = 'encoder{:02d}/transformer_encode/src_self_attn/multi_head_attention/q_weight'.format(
layer_id)
if 'key/bias' in key:
key = 'encoder{:02d}/transformer_encode/src_self_attn/multi_head_attention/k_bias'.format(
layer_id)
if 'key/kernel' in key:
key = 'encoder{:02d}/transformer_encode/src_self_attn/multi_head_attention/k_weight'.format(
layer_id)
if 'value/bias' in key:
key = 'encoder{:02d}/transformer_encode/src_self_attn/multi_head_attention/v_bias'.format(
layer_id)
if 'value/kernel' in key:
key = 'encoder{:02d}/transformer_encode/src_self_attn/multi_head_attention/v_weight'.format(
layer_id)
if 'attention/output/LayerNorm/beta' in key:
key = 'encoder{:02d}/transformer_encode/enc_layer_norm1/layer_normalization/beta'.format(
layer_id)
weight = numpy.reshape(weight, (1, 1, 768))
elif 'output/LayerNorm/beta' in key:
key = 'encoder{:02d}/transformer_encode/enc_layer_norm2/layer_normalization/beta'.format(
layer_id)
weight = numpy.reshape(weight, (1, 1, 768))
if 'attention/output/LayerNorm/gamma' in key:
key = 'encoder{:02d}/transformer_encode/enc_layer_norm1/layer_normalization/gamma'.format(
layer_id)
weight = numpy.reshape(weight, (1, 1, 768))
elif 'output/LayerNorm/gamma' in key:
key = 'encoder{:02d}/transformer_encode/enc_layer_norm2/layer_normalization/gamma'.format(
layer_id)
weight = numpy.reshape(weight, (1, 1, 768))
if 'attention/output/dense/bias' in key:
key = 'encoder{:02d}/transformer_encode/src_self_attn/multi_head_attention/out_bias'.format(
layer_id)
elif 'output/dense/bias' in key:
key = 'encoder{:02d}/transformer_encode/enc_affine2/affine/b'.format(
layer_id)
if 'intermediate/dense/bias' in key:
key = 'encoder{:02d}/transformer_encode/enc_affine1/affine/b'.format(
layer_id)
if 'attention/output/dense/kernel' in key:
key = 'encoder{:02d}/transformer_encode/src_self_attn/multi_head_attention/out_weight'.format(
layer_id)
elif 'output/dense/kernel' in key:
key = 'encoder{:02d}/transformer_encode/enc_affine2/affine/W'.format(
layer_id)
if 'intermediate/dense/kernel' in key:
key = 'encoder{:02d}/transformer_encode/enc_affine1/affine/W'.format(
layer_id)
if 'embeddings/LayerNorm/' in key:
key = key.replace('bert/embeddings/LayerNorm',
'embed/layer_normalization')
weight = numpy.reshape(weight, (1, 1, 768))
if 'word_embeddings' in key:
key = 'word_embeddings/embed/W'
if 'token_type_embeddings' in key:
key = 'token_type_embeddings/embed/W'
if 'position_embeddings' in key:
key = 'position_embeddings/embed/W'
if 'pooler/dense/bias' in key:
key = 'pooler/affine/b'
if 'pooler/dense/kernel' in key:
key = 'pooler/affine/W'
if 'seq_relationship/output_weights' in key:
key = 'affine_seq_class/affine/W'
weight = numpy.transpose(weight)
if 'seq_relationship/output_bias' in key:
key = 'affine_seq_class/affine/b'
params = PF.get_parameter_or_create(key, shape=weight.shape)
params.d = weight
nn.parameter.save_parameters(h5_file)