Secure your code as it's written. Use Snyk Code to scan source code in minutes - no build needed - and fix issues immediately.
def testSKLearnGraph(self):
sklearn_model, prototype = get_sklearn_model_and_prototype()
path = f'{time.time()}.onnx'
self.assertRaises(TypeError, save_model, sklearn_model, path)
save_model(sklearn_model, path, prototype=prototype)
model = load_model(path)
os.remove(path)
con = self.get_client()
con.modelset('onnx_skl_model', Backend.onnx, Device.cpu, model)
con.tensorset('a', Tensor.scalar(DType.float, *([1] * 13)))
con.modelrun('onnx_skl_model', ['a'], ['outfromonnxskl'])
tensor = con.tensorget('outfromonnxskl')
self.assertEqual(len(tensor.value), 1)
def testTFGraph(self):
_ = get_tf_graph()
init = tf.global_variables_initializer()
sess = tf.Session()
sess.run(init)
path = f'{time.time()}.pb'
save_model(sess, path, output=['output'])
model = load_model(path)
os.remove(path)
con = self.get_client()
con.modelset(
'tfmodel', Backend.tf, Device.cpu, model,
input=['input'], output=['output'])
con.tensorset('a', Tensor.scalar(DType.float, 2))
con.modelrun('tfmodel', ['a'], 'c')
tensor = con.tensorget('c')
self.assertEqual([13], tensor.value)
def testONNXGraph(self):
onnx_model = get_onnx_model()
path = f'{time.time()}.onnx'
save_model(onnx_model, path)
model = load_model(path)
os.remove(path)
con = self.get_client()
con.modelset('onnxmodel', Backend.onnx, Device.cpu, model)
con.tensorset('a', Tensor.scalar(DType.float, 2, -1))
con.modelrun('onnxmodel', ['a'], ['c'])
tensor = con.tensorget('c')
self.assertEqual([2.0, 0.0], tensor.value)
def testPyTorchGraph(self):
torch_graph = MyModule()
path = f'{time.time()}.pt'
save_model(torch_graph, path)
model = load_model(path)
os.remove(path)
con = self.get_client()
con.modelset('ptmodel', Backend.torch, Device.cpu, model)
con.tensorset('a', Tensor.scalar(DType.float, 2, 5))
con.tensorset('b', Tensor.scalar(DType.float, 3, 7))
con.modelrun('ptmodel', ['a', 'b'], 'c')
tensor = con.tensorget('c')
self.assertEqual([5, 12], tensor.value)