def convert(infile, outdir, **kwargs):
"""Convert pb.
Args:
infile: Input path.
outdir: Output path.
**kwargs: Other args for converting.
Returns:
None.
"""
logging_level = kwargs.get("logging_level", "INFO")
ext_data_dir = kwargs.get("extdatadir")
common.logger.setLevel(logging_level)
common.logger.handlers[0].setLevel(logging_level)
common.logger.info("Start converting onnx pb to tf saved model")
# load external data if the file directory is provided
if ext_data_dir:
onnx_model = onnx.load(infile, load_external_data=False)
load_external_data_for_model(onnx_model, ext_data_dir)
else:
onnx_model = onnx.load(infile)
tf_rep = backend.prepare(onnx_model, **kwargs)
tf_rep.export_graph(outdir)
common.logger.info("Converting completes successfully.")
def load_model(f: Union[IO[bytes], Text],
format: Optional[Any] = None,
load_external_data: bool = True) -> ModelProto:
'''
Loads a serialized ModelProto into memory
load_external_data is true if the external data under the same directory of the model and load the external data
If not, users need to call load_external_data_for_model with directory to load
Arguments:
f: can be a file-like object (has "read" function) or a string containing a file name
format: for future use
Returns:
Loaded in-memory ModelProto
'''
s = _load_bytes(f)
model = load_model_from_string(s, format=format)
if load_external_data:
model_filepath = _get_file_path(f)
if model_filepath:
base_dir = os.path.dirname(model_filepath)
load_external_data_for_model(model, base_dir)
return model
def load_model(
f,
format=None,
load_external_data=True
): # type: (Union[IO[bytes], Text], Optional[Any], bool) -> ModelProto
'''
Loads a serialized ModelProto into memory
@params
f can be a file-like object (has "read" function) or a string containing a file name
format is for future use
@return
Loaded in-memory ModelProto
'''
s = _load_bytes(f)
model = load_model_from_string(s, format=format)
if load_external_data:
model_filepath = _get_file_path(f)
if model_filepath:
base_dir = os.path.dirname(model_filepath)
load_external_data_for_model(model, base_dir)
return model
def test_load_external_data_for_model(self): # type: () -> None
model = onnx.load_model(self.model_filename, load_external_data=False)
load_external_data_for_model(model, self.temp_dir)
initializer_tensor = model.graph.initializer[0]
self.assertTrue(np.allclose(to_array(initializer_tensor), self.initializer_value))
attribute_tensor = model.graph.node[0].attribute[0].t
self.assertTrue(np.allclose(to_array(attribute_tensor), self.attribute_value))
def load_initializers_from_onnx(model_path, load_optimizer=False):
"""Load initial weights from an onnx checkpoint.
Args:
model_path (str): Path to onnx file.
Returns:
Dict: Mapping of popart weight names to numpy values.
"""
initializers = {}
# By default onnx.load will look for initializers in the same dir as onnx model.
# However builder.saveIntializersExternally takes real path or path relative to run dir
# and stores it in the onnxproto.
model = onnx.load(model_path, load_external_data=False)
has_external_data = any(
is_external_weight(weight) for weight in model.graph.initializer)
if has_external_data:
load_external_data_for_model(model, '')
optimizer_prefix = (popart.reservedAccl1Prefix(),
popart.reservedAccl2Prefix(),
popart.reservedAcclPrefix(),
popart.reservedAccumPrefix(),
popart.reservedStepPrefix())
for weight in model.graph.initializer:
is_optimizer_state = any(x in weight.name for x in optimizer_prefix)
if not load_optimizer and is_optimizer_state:
continue
if is_external_weight(
weight) or weight.data_type != onnx.TensorProto.FLOAT16:
np_weight = numpy_helper.to_array(weight)
else:
int_data = np.asarray(weight.int32_data, np.int32)
np_weight = int_data.view(dtype=np.float16).reshape(weight.dims)
if is_optimizer_state:
initializers[weight.name] = np_weight.astype(np.float32)
else:
initializers[weight.name] = np_weight
initializers = handle_split_qkv(initializers)
initializers = handle_split_word_embedding(initializers)
return initializers
def test_check_model_by_model(self): # type: () -> None
model = onnx.load_model(self.model_filename, load_external_data=False)
load_external_data_for_model(model,
self.temp_dir) # Exceeds maximum protobuf
self.assertRaises(ValueError, checker.check_model,
model) # checker catches 2GB models and throw error
def test_check_model_by_model(self): # type: () -> None
model = onnx.load_model(self.model_filename, load_external_data=False)
with pytest.raises(ValueError):
load_external_data_for_model(
model, self.temp_dir) # Exceeds maximum protobuf
checker.check_model(model) # checker catches 2GB models as well
