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