def assertONNX(self, f, args, params=tuple(), **kwargs):
if isinstance(f, nn.Module):
m = f
else:
m = FuncModule(f, params)
onnx_model_pb = export_to_string(m, args, **kwargs)
model_def = self.assertONNXExpected(onnx_model_pb)
if _onnx_test:
test_function = inspect.stack()[1][0].f_code.co_name
test_name = test_function[0:4] + "_operator" + test_function[4:]
output_dir = os.path.join(test_onnx_common.pytorch_operator_dir, test_name)
# Assume:
# 1) the old test should be delete before the test.
# 2) only one assertONNX in each test, otherwise will override the data.
assert not os.path.exists(output_dir), "{} should not exist!".format(output_dir)
os.makedirs(output_dir)
with open(os.path.join(output_dir, "model.onnx"), 'wb') as file:
file.write(model_def.SerializeToString())
data_dir = os.path.join(output_dir, "test_data_set_0")
os.makedirs(data_dir)
if isinstance(args, Variable):
args = (args,)
for index, var in enumerate(flatten(args)):
tensor = numpy_helper.from_array(var.data.numpy())
with open(os.path.join(data_dir, "input_{}.pb".format(index)), 'wb') as file:
file.write(tensor.SerializeToString())
outputs = m(*args)
if isinstance(outputs, Variable):
outputs = (outputs,)
for index, var in enumerate(flatten(outputs)):
tensor = numpy_helper.from_array(var.data.numpy())
with open(os.path.join(data_dir, "output_{}.pb".format(index)), 'wb') as file:
file.write(tensor.SerializeToString())
def assertONNX(self, f, args, params=tuple(), **kwargs):
if isinstance(f, nn.Module):
m = f
else:
m = FuncModule(f, params)
onnx_model_pb = export_to_string(m, args, **kwargs)
model_def = self.assertONNXExpected(onnx_model_pb)
if _onnx_test:
test_function = inspect.stack()[1][0].f_code.co_name
test_name = test_function[0:4] + "_operator" + test_function[4:]
output_dir = os.path.join(test_onnx_common.pytorch_operator_dir, test_name)
# Assume:
# 1) the old test should be delete before the test.
# 2) only one assertONNX in each test, otherwise will override the data.
assert not os.path.exists(output_dir), "{} should not exist!".format(output_dir)
os.makedirs(output_dir)
with open(os.path.join(output_dir, "model.onnx"), 'wb') as file:
file.write(model_def.SerializeToString())
data_dir = os.path.join(output_dir, "test_data_set_0")
os.makedirs(data_dir)
if isinstance(args, Variable):
args = (args,)
for index, var in enumerate(flatten(args)):
tensor = numpy_helper.from_array(var.data.numpy())
with open(os.path.join(data_dir, "input_{}.pb".format(index)), 'wb') as file:
file.write(tensor.SerializeToString())
outputs = m(*args)
if isinstance(outputs, Variable):
outputs = (outputs,)
for index, var in enumerate(flatten(outputs)):
tensor = numpy_helper.from_array(var.data.numpy())
with open(os.path.join(data_dir, "output_{}.pb".format(index)), 'wb') as file:
file.write(tensor.SerializeToString())
def assertONNX(self, f, args, params=None, **kwargs):
if params is None:
params = ()
if isinstance(f, nn.Module):
m = f
else:
m = FuncModule(f, params)
m.eval()
onnx_model_pbtxt = export_to_pbtxt(m, args, **kwargs)
subname = kwargs.pop("subname", None)
self.assertExpected(onnx_model_pbtxt, subname)
if _onnx_dep:
onnx_model_pb = export_to_pb(m, args, **kwargs)
import onnx
import onnx.checker
import onnx.numpy_helper
import test_onnx_common
model_def = onnx.ModelProto.FromString(onnx_model_pb)
onnx.checker.check_model(model_def)
if _onnx_test:
test_function = inspect.stack()[1][0].f_code.co_name
test_name = test_function[0:4] + "_operator" + test_function[4:]
output_dir = os.path.join(
test_onnx_common.pytorch_operator_dir, test_name)
# Assume:
# 1) the old test should be delete before the test.
# 2) only one assertONNX in each test, otherwise will override the data.
assert not os.path.exists(
output_dir), "{} should not exist!".format(output_dir)
os.makedirs(output_dir)
with open(os.path.join(output_dir, "model.onnx"),
"wb") as file:
file.write(model_def.SerializeToString())
data_dir = os.path.join(output_dir, "test_data_set_0")
os.makedirs(data_dir)
if isinstance(args, Variable):
args = (args, )
for index, var in enumerate(flatten(args)):
tensor = onnx.numpy_helper.from_array(var.data.numpy())
with open(os.path.join(data_dir, f"input_{index}.pb"),
"wb") as file:
file.write(tensor.SerializeToString())
outputs = m(*args)
if isinstance(outputs, Variable):
outputs = (outputs, )
for index, var in enumerate(flatten(outputs)):
tensor = onnx.numpy_helper.from_array(var.data.numpy())
with open(os.path.join(data_dir, f"output_{index}.pb"),
"wb") as file:
file.write(tensor.SerializeToString())
def run_embed_params(proto, model, input, state_dict=None, use_gpu=True):
"""
This is only a helper debug function so we can test embed_params=False
case as well on pytorch front
This should likely be removed from the release version of the code
"""
device = "CPU"
if use_gpu:
device = "CUDA"
model_def = onnx.ModelProto.FromString(proto)
onnx.checker.check_model(model_def)
prepared = c2.prepare(model_def, device=device)
if state_dict:
parameters = []
# Passed in state_dict may have a different order. Make
# sure our order is consistent with the model's order.
# TODO: Even better: keyword arguments!
for k in model.state_dict():
if k in state_dict:
parameters.append(state_dict[k])
else:
parameters = list(model.state_dict().values())
W = {}
for k, v in zip(model_def.graph.input, flatten((input, parameters))):
if isinstance(v, Variable):
W[k.name] = v.data.cpu().numpy()
else:
W[k.name] = v.cpu().numpy()
caffe2_out = prepared.run(inputs=W)
return caffe2_out
def run_embed_params(proto, model, input, state_dict=None, use_gpu=True):
"""
This is only a helper debug function so we can test embed_params=False
case as well on pytorch front
This should likely be removed from the release version of the code
"""
device = 'CPU'
if use_gpu:
device = 'CUDA'
model_def = onnx.ModelProto.FromString(proto)
onnx.checker.check_model(model_def)
prepared = c2.prepare(model_def, device=device)
if state_dict:
parameters = []
# Passed in state_dict may have a different order. Make
# sure our order is consistent with the model's order.
# TODO: Even better: keyword arguments!
for k in model.state_dict():
if k not in state_dict:
# Once PyTorch Module adds unnecessary paramter, the old pre-trained model does not have it.
# Just simply pass the new one.
# TODO: Please don't export unnecessary parameter.
parameters.append(model.state_dict()[k])
else:
parameters.append(state_dict[k])
else:
parameters = list(model.state_dict().values())
W = {}
for k, v in zip(model_def.graph.input, flatten((input, parameters))):
if isinstance(v, Variable):
W[k.name] = v.data.cpu().numpy()
else:
W[k.name] = v.cpu().numpy()
caffe2_out = prepared.run(inputs=W)
return caffe2_out
