wehub-resource-sync
123 lines
4.4 KiB
Python
123 lines
4.4 KiB
Python
#!/usr/bin/python
|
|
import os
|
|
import sys
|
|
import numpy as np
|
|
import tensorflow as tf
|
|
|
|
def makeDirForPath(filename):
|
|
if filename.find('/') < 0:
|
|
return
|
|
names = filename.split('/')
|
|
dirname = ""
|
|
for l in range(0, len(names)-1):
|
|
dirname = dirname + names[l] + '/'
|
|
print(dirname)
|
|
if os.path.exists(dirname):
|
|
return
|
|
os.makedirs(dirname)
|
|
|
|
class TestModel():
|
|
def __copy_to_here(self, modelName):
|
|
newModel = 'tf/test.pb'
|
|
print(os.popen("mkdir tf").read())
|
|
print(os.popen("cp " + modelName + ' ' + newModel).read())
|
|
self.modelName = newModel
|
|
self.model = self.__load_graph(self.modelName)
|
|
self.inputOps, self.outputOps = self.__analyze_inputs_outputs(self.model)
|
|
self.outputs = [output.name for output in self.outputOps]
|
|
def __init__(self, modelName):
|
|
self.__copy_to_here(modelName)
|
|
def __run_mnn(self):
|
|
mnnconvert_name = 'MNNConvert.exe' if os.name == 'nt' else './MNNConvert'
|
|
if not os.path.exists(mnnconvert_name):
|
|
print("./MNNConvert not exist in this path. Use pymnn instead of C++ to test")
|
|
mnnconvert_name = 'mnnconvert'
|
|
convert = mnnconvert_name + ' -f TF --bizCode MNN --modelFile tf/test.pb --MNNModel convert_cache.mnn --keepInputFormat=1 --testdir tf'
|
|
result = os.popen(convert).read()
|
|
print(result)
|
|
return result
|
|
def __load_graph(self, filename):
|
|
f = tf.io.gfile.GFile(filename, "rb")
|
|
graph_def = tf.compat.v1.GraphDef()
|
|
graph_def.ParseFromString(f.read())
|
|
tf.import_graph_def(graph_def, name='')
|
|
graph = tf.compat.v1.get_default_graph()
|
|
return graph
|
|
def __analyze_inputs_outputs(self, graph):
|
|
ops = graph.get_operations()
|
|
outputs_set = set(ops)
|
|
inputs = []
|
|
testop = None
|
|
for op in ops:
|
|
if op.name == specifyOpName:
|
|
testop = op
|
|
if len(op.inputs) == 0 and op.type != 'Const':
|
|
inputs.append(op)
|
|
else:
|
|
for input_tensor in op.inputs:
|
|
if input_tensor.op in outputs_set:
|
|
outputs_set.remove(input_tensor.op)
|
|
outputs = [op for op in outputs_set if op.type != 'Assert' and op.type != 'Const']
|
|
if testop != None:
|
|
outputs = [ testop ]
|
|
return (inputs, outputs)
|
|
def __get_shape(self, op):
|
|
shape = [s.value if tf.__version__[0] == '1' else s for s in op.outputs[0].shape]
|
|
for i in range(len(shape)):
|
|
if shape[i] == None:
|
|
shape[i] = 1
|
|
return shape
|
|
def __run_tf(self):
|
|
jsonDict = {}
|
|
jsonDict['inputs'] = []
|
|
jsonDict['outputs'] = []
|
|
inputs = {}
|
|
print(self.modelName)
|
|
for inputVar in self.inputOps:
|
|
inp = {}
|
|
inp['name'] = inputVar.name
|
|
inp['shape'] = self.__get_shape(inputVar)
|
|
inputs[inputVar.name + ':0'] = np.random.uniform(0.1, 1.2, inp['shape']).astype(np.sctypeDict[inputVar.outputs[0].dtype.name])
|
|
jsonDict['inputs'].append(inp)
|
|
print([output.name for output in self.outputOps])
|
|
for output in self.outputOps:
|
|
jsonDict['outputs'].append(output.name)
|
|
|
|
import json
|
|
jsonString = json.dumps(jsonDict, indent=4)
|
|
with open('tf/input.json', 'w') as f:
|
|
f.write(jsonString)
|
|
|
|
print('inputs:')
|
|
for key in inputs:
|
|
print(key)
|
|
name = 'tf/' + key[:-2] + '.txt'
|
|
makeDirForPath(name)
|
|
f = open(name, 'w')
|
|
np.savetxt(f, inputs[key].flatten())
|
|
f.close()
|
|
sess = tf.compat.v1.Session()
|
|
outputs_tensor = [(output + ':0') for output in self.outputs]
|
|
outputs = sess.run(outputs_tensor, inputs)
|
|
print('outputs:')
|
|
for i in range(len(outputs)):
|
|
outputName = self.outputs[i]
|
|
name = 'tf/' + outputName + '.txt'
|
|
makeDirForPath(name)
|
|
# print(name, outputs[i].shape)
|
|
f = open(name, 'w')
|
|
np.savetxt(f, outputs[i].flatten())
|
|
f.close()
|
|
def Test(self):
|
|
self.__run_tf()
|
|
res = self.__run_mnn()
|
|
return res
|
|
|
|
if __name__ == '__main__':
|
|
modelName = sys.argv[1]
|
|
specifyOpName = None
|
|
if len(sys.argv) > 2:
|
|
specifyOpName = sys.argv[2]
|
|
t = TestModel(modelName)
|
|
t.Test()
|