Francesco Gatti
134 lines
3.9 KiB
Python
134 lines
3.9 KiB
Python
import keras
|
|
from keras.models import load_model
|
|
import keras.backend.tensorflow_backend as KTF
|
|
import numpy as np
|
|
import argparse
|
|
import tensorflow as tf
|
|
import os
|
|
import msgpack
|
|
import lmdb
|
|
import random
|
|
|
|
def export_dense(name, weights, bias):
|
|
print "######## EXPORT", name, "LAYER ########"
|
|
print "Original weighs:"
|
|
print weights
|
|
print bias, "\n"
|
|
|
|
#input, filters
|
|
I, C = np.shape(weights)
|
|
B = np.shape(bias)
|
|
print "w shape: ", I, C
|
|
print "b shape: ", B
|
|
|
|
wgs = [ [ j[i] for j in weights ] for i in xrange(C) ]
|
|
wgs = np.array(wgs, dtype=np.float32)
|
|
|
|
print "REPOSITIONED WEIGHTS:"
|
|
print wgs
|
|
|
|
bias = np.array(bias, dtype=np.float32)
|
|
wgs.tofile(name + ".bin", format="f")
|
|
bias.tofile(name + ".bias.bin", format="f")
|
|
print "WEIGHTS saved\n"
|
|
|
|
def export_conv2d(name, weights, bias):
|
|
print "######## EXPORT", name, "LAYER ########"
|
|
print "Original weighs:"
|
|
print weights
|
|
print bias, "\n"
|
|
|
|
# height, width, input, filters
|
|
H, W, N, C = np.shape(weights)
|
|
B = np.shape(bias)
|
|
print "w shape: ", N, C, H, W
|
|
print "b shape: ", B
|
|
|
|
wgs = weights.transpose()
|
|
wgs = wgs.transpose(0, 1, 3, 2)
|
|
print "Final shape:", np.shape(wgs)
|
|
wgs = np.array(wgs.flatten(), dtype=np.float32)
|
|
|
|
print "REPOSITIONED WEIGHTS:"
|
|
print wgs
|
|
|
|
bias = np.array(bias, dtype=np.float32)
|
|
|
|
wgs.tofile(name + ".bin", format="f")
|
|
bias.tofile(name + ".bias.bin", format="f")
|
|
print "WEIGHTS saved\n"
|
|
|
|
def export_conv3d(name, weights, bias):
|
|
print "######## EXPORT", name, "LAYER ########"
|
|
print "Original weighs:"
|
|
print weights
|
|
print bias, "\n"
|
|
|
|
print np.shape(weights)
|
|
# height, width, input, thickness, filters
|
|
H, W, T, N, C = np.shape(weights)
|
|
B = np.shape(bias)
|
|
print "w shape: ", T, C, H, W #thickness is number of images for cudnn
|
|
print "b shape: ", B
|
|
|
|
wgs = weights.transpose()
|
|
wgs = wgs.transpose(0, 1, 4, 3, 2)
|
|
print "Final shape:", np.shape(wgs)
|
|
wgs = np.array(wgs.flatten(), dtype=np.float32)
|
|
|
|
print "REPOSITIONED WEIGHTS:"
|
|
print wgs
|
|
|
|
bias = np.array(bias, dtype=np.float32)
|
|
|
|
wgs.tofile(name + ".bin", format="f")
|
|
bias.tofile(name + ".bias.bin", format="f")
|
|
print "WEIGHTS saved\n"
|
|
|
|
|
|
def get_session(gpu_fraction=0.5):
|
|
gpu_options = tf.GPUOptions(allow_growth=True)
|
|
#per_process_gpu_memory_fraction=gpu_fraction)
|
|
return tf.Session(config=tf.ConfigProto(gpu_options=gpu_options))
|
|
|
|
|
|
#https://github.com/fchollet/keras/wiki/Converting-convolution-kernels-from-Theano-to-TensorFlow-and-vice-versa
|
|
if __name__ == '__main__':
|
|
KTF.set_session(get_session())
|
|
|
|
parser = argparse.ArgumentParser(description='KERAS WEIGHTS EXPORTER TO CUDNN')
|
|
parser.add_argument('model',type=str,
|
|
help='Path to model h5 file. Model should be on the same path.')
|
|
parser.add_argument('--output', type=str, help="output directory", default="layers")
|
|
|
|
args = parser.parse_args()
|
|
|
|
print "DATA FORMAT: ", keras.backend.image_data_format()
|
|
|
|
print "Load model: ", args.model
|
|
model = load_model(args.model)
|
|
weights = model.get_weights()
|
|
|
|
ws = np.shape(weights)
|
|
print "Weights shape:", ws
|
|
|
|
if not os.path.exists(args.output):
|
|
os.makedirs(args.output)
|
|
|
|
num = 0
|
|
name_num = 0
|
|
for l in model.layers:
|
|
name = l.name
|
|
if name.startswith("conv3d"):
|
|
export_conv3d(args.output + "/conv" + str(name_num), weights[num], weights[num+1])
|
|
elif name.startswith("conv2d"):
|
|
export_conv2d(args.output + "/conv" + str(name_num), weights[num], weights[num+1])
|
|
elif name.startswith("dense"):
|
|
export_dense(args.output + "/dense" + str(name_num), weights[num], weights[num+1])
|
|
else:
|
|
print "skip:", name, "has no weights"
|
|
continue
|
|
name_num += 1
|
|
num += 2
|
|
|