def save_framework_info(framework: Any, path: Optional[str] = None):
"""
Save the framework info for a given framework.
If path is provided, will save to a json file at that path.
If path is not provided, will print out the info.
:param framework: The item to detect the ML framework for.
See :func:`detect_framework` for more information.
:type framework: Any
:param path: The path, if any, to save the info to in json format.
If not provided will print out the info.
:type path: Optional[str]
"""
_LOGGER.debug(
"saving framework info for framework %s to %s",
framework,
path if path else "sys.out",
)
info = (framework_info(framework)
if not isinstance(framework, FrameworkInfo) else framework)
if path:
path = clean_path(path)
create_parent_dirs(path)
with open(path, "w") as file:
file.write(info.json())
_LOGGER.info("saved framework info for framework %s in file at %s",
framework, path),
else:
print(info.json(indent=4))
_LOGGER.info("printed out framework info for framework %s", framework)
def __init__(
self,
model: keras.Model,
output_dir: str,
):
self._model = model
self._output_dir = clean_path(output_dir)
def __init__(
self,
root: str = default_dataset_path(
"imagenet"), # default to imagenet location
train: bool = True,
rand_trans: bool = False,
image_size: int = 224,
):
if torchvision_import_error is not None:
raise torchvision_import_error
root = clean_path(root)
non_rand_resize_scale = 256.0 / 224.0 # standard used
init_trans = ([
transforms.RandomResizedCrop(image_size),
transforms.RandomHorizontalFlip(),
] if rand_trans else [
transforms.Resize(round(non_rand_resize_scale * image_size)),
transforms.CenterCrop(image_size),
])
trans = [
*init_trans,
transforms.ToTensor(),
transforms.Normalize(mean=IMAGENET_RGB_MEANS,
std=IMAGENET_RGB_STDS),
]
root = os.path.join(os.path.abspath(os.path.expanduser(root)),
"train" if train else "val")
super().__init__(root, transform=transforms.Compose(trans))
if train:
# make sure we dont preserve the folder structure class order
random.shuffle(self.samples)
def __init__(self, download_root: str, dataset_size: ImagenetteSize,
download: bool):
self._download_root = clean_path(download_root)
self._dataset_size = dataset_size
self._download = download
if dataset_size == ImagenetteSize.s160:
self._extract_name = "imagenette-160"
elif dataset_size == ImagenetteSize.s320:
self._extract_name = "imagenette-320"
elif dataset_size == ImagenetteSize.full:
self._extract_name = "imagenette"
else:
raise ValueError(
"Unknown ImagenetteSize given of {}".format(dataset_size))
self._extracted_root = os.path.join(self._download_root,
self._extract_name)
if download:
self._download_and_extract()
else:
file_path = "{}.tar".format(self._extracted_root)
if not os.path.exists(file_path):
raise ValueError(
"could not find original tar for the dataset at {}".format(
file_path))
def __init__(
self,
root: str = default_dataset_path("imagenet"),
train: bool = True,
rand_trans: bool = False,
image_size: Union[None, int, Tuple[int, int]] = 224,
pre_resize_transforms=SplitsTransforms(
train=(
random_scaling_crop(),
tf.image.random_flip_left_right,
),
val=(imagenet_pre_resize_processor(), ),
),
post_resize_transforms=SplitsTransforms(
train=(torch_imagenet_normalizer(), ),
val=(torch_imagenet_normalizer(), )),
):
root = clean_path(root)
super().__init__(
root,
train,
image_size=image_size,
pre_resize_transforms=pre_resize_transforms,
post_resize_transforms=post_resize_transforms,
)
if train:
# make sure we don't preserve the folder structure class order
random.shuffle(self.samples)
def save_benchmark_results(
model: Any,
data: Any,
batch_size: int,
iterations: int,
warmup_iterations: int,
framework: Optional[str],
provider: Optional[str] = None,
device: Optional[str] = None,
save_path: Optional[str] = None,
framework_args: Dict[str, Any] = {},
show_progress: bool = False,
):
"""
Saves the benchmark results ran for specific framework.
If path is provided, will save to a json file at the path.
If path is not provided, will print out the info.
If no framework is provided, will detect the framework based on the model.
:param model: model to benchmark
:param data: data to benchmark
:param batch_size: batch size
:param iterations: number of iterations
:param warmup_iterations: number of warmup iterations
:param framework: the specific framework run the benchmark in
:param provider: the specific inference provider to use
:param device: the specific device to use
:param save_path: path to save the benchmark results
:param framework_args: additional framework specific arguments to
pass to the runner
:param show_progress: True to show a tqdm bar when running, False otherwise
"""
results = execute_in_sparseml_framework(
framework if framework is not None else model,
"run_benchmark",
model,
data,
batch_size=batch_size,
iterations=iterations,
warmup_iterations=warmup_iterations,
provider=provider,
device=device,
framework_args=framework_args,
show_progress=show_progress,
)
if save_path:
save_path = clean_path(save_path)
create_parent_dirs(save_path)
with open(save_path, "w") as file:
file.write(results.json(indent=4))
_LOGGER.info(f"saved benchmark results in file at {save_path}"),
else:
print(results.json(indent=4))
_LOGGER.info("printed out benchmark results")
def load(file_path) -> "ModelInfo":
"""
:param file_path: file path to JSON file to load ModelInfo object from
:return: the loaded ModelInfo object
"""
file_path = clean_path(file_path)
with open(file_path, "r") as file:
model_info_dict = json.load(file)
return ModelInfo.from_dict(model_info_dict)
def save(self, file_path: str):
"""
:param file_path: the file path to save the yaml config representation to
"""
file_path = clean_path(file_path)
create_parent_dirs(file_path)
with open(file_path, "w") as yaml_file:
yaml_file.write(str(self))
def __init__(
self,
module: Module,
output_dir: str,
):
if is_parallel_model(module):
module = module.module
self._module = deepcopy(module).to("cpu").eval()
self._output_dir = clean_path(output_dir)
def load_json(path: str):
"""
:param path: the path to load a previous analysis from
:return: the ModelAnalyzer instance from the json
"""
path = clean_path(path)
with open(path, "r") as file:
objs = json.load(file)
return ModelAnalyzer.from_dict(objs)
def save_descs(descs: List, path: str):
"""
Save a list of AnalyzedLayerDesc to a json file
:param descs: a list of descriptions to save
:param path: the path to save the descriptions at
"""
path = clean_path(path)
create_parent_dirs(path)
save_obj = {"descriptions": [desc.dict() for desc in descs]}
with open(path, "w") as file:
json.dump(save_obj, file)
def save_json(self, path: str):
"""
:param path: the path to save the json file at representing the analyzed
results
"""
if not path.endswith(".json"):
path += ".json"
path = clean_path(path)
create_parent_dirs(path)
with open(path, "w") as file:
dictionary = self.dict()
json.dump(dictionary, file, indent=2)
def check_load_model(model: Union[str, ModelProto]) -> ModelProto:
"""
Load an ONNX model from a given file path if supplied.
If already a model proto, then returns.
:param model: the model proto or path to the model onnx file to check for loading
:return: the loaded onnx ModelProto
"""
if isinstance(model, ModelProto):
return model
if isinstance(model, str):
return onnx.load(clean_path(model))
raise ValueError("unknown type given for model: {}".format(model))
def load_benchmark_info(load: str) -> BenchmarkInfo:
"""
Load the benchmark info from a file or raw json.
If load exists as a path, will read from the file and use that.
Otherwise will try to parse the input as a raw json str.
:param load: Either a file path to a json file or a raw json string.
:type load: str
:return: The loaded benchmark info.
:rtype: FrameworkInfo
"""
loaded_path = clean_path(load)
if os.path.exists(loaded_path):
with open(loaded_path, "r") as file:
load = file.read()
info = BenchmarkInfo.parse_raw(load)
return info
def train_setup():
def _create_optim(_model: Module) -> Optimizer:
return SGD(
_model.parameters(),
lr=0.1,
momentum=0.9,
nesterov=True,
weight_decay=0.0001,
)
# fill in the appropriate values below for your training flow
train(
working_dir=clean_path("."),
config_path="/PATH/TO/CONFIG.yaml",
model=None,
train_dataset=None,
val_dataset=None,
batch_size=64,
optim_const=_create_optim,
loss=None,
devices=default_device(),
)
def __init__(
self,
root: str,
train: bool,
image_size: Union[None, int, Tuple[int, int]] = 224,
pre_resize_transforms: Union[
SplitsTransforms, None] = SplitsTransforms(
train=(
random_scaling_crop(),
tensorflow.image.random_flip_left_right,
),
val=None,
),
post_resize_transforms: Union[SplitsTransforms,
None] = SplitsTransforms(
train=(
default_imagenet_normalizer(), ),
val=(
default_imagenet_normalizer(), ),
),
):
self._root = os.path.join(clean_path(root),
"train" if train else "val")
if not os.path.exists(self._root):
raise ValueError("Data set folder {} must exist".format(
self._root))
self._train = train
if image_size is not None:
self._image_size = (image_size if isinstance(image_size, tuple)
else (image_size, image_size))
else:
self._image_size = None
self._pre_resize_transforms = pre_resize_transforms
self._post_resize_transforms = post_resize_transforms
self._num_images = len(
[None for _ in glob.glob(os.path.join(self._root, "*", "*"))])
self._num_classes = len(
[None for _ in glob.glob(os.path.join(self._root, "*", ""))])
def load_descs(path: str) -> List:
"""
Load a list of AnalyzedLayerDesc from a json file
:param path: the path to load the descriptions from
:return: the loaded list of AnalyzedLayerDesc
"""
path = clean_path(path)
with open(path, "r") as file:
obj = json.load(file)
descs = []
for desc_obj in obj["descriptions"]:
desc_obj["type_"] = desc_obj["type"]
del desc_obj["type"]
del desc_obj["terminal"]
del desc_obj["prunable"]
descs.append(AnalyzedLayerDesc(**desc_obj))
return descs
def __init__(
self,
root: str,
train: bool,
image_size: int = 224,
pre_resize_transforms: Union[
SplitsTransforms, None] = SplitsTransforms(
train=(
random_scaling_crop(),
tf_compat.image.random_flip_left_right,
tf_compat.image.random_flip_up_down,
),
val=None,
),
post_resize_transforms: Union[SplitsTransforms,
None] = SplitsTransforms(
train=(imagenet_normalizer, ),
val=(
center_square_crop(),
imagenet_normalizer,
),
),
):
self._root = os.path.join(clean_path(root),
"train" if train else "val")
if not os.path.exists(self._root):
raise ValueError("Data set folder {} must exist".format(
self._root))
self._train = train
self._image_size = image_size
self._pre_resize_transforms = pre_resize_transforms
self._post_resize_transforms = post_resize_transforms
self._num_images = len(
[None for _ in glob.glob(os.path.join(self._root, "*", "*"))])
self._num_classes = len(
[None for _ in glob.glob(os.path.join(self._root, "*", ""))])
def __init__(self, output_dir: str):
self._output_dir = clean_path(output_dir)
def train(
working_dir: str,
config_path: str,
model: Module,
train_dataset: Dataset,
val_dataset: Dataset,
batch_size: int,
optim_const: Callable[[Module], Optimizer],
loss: Union[LossWrapper, Callable[[Any, Any], Tensor]],
devices: str,
):
"""
Dataset setup
"""
LOGGER.info("batch_size set to {}".format(batch_size))
LOGGER.info("train_dataset set to {}".format(train_dataset))
LOGGER.info("val_dataset set to {}".format(val_dataset))
train_loader = DataLoader(
train_dataset,
batch_size=batch_size,
shuffle=True,
num_workers=8,
pin_memory=True,
)
val_loader = DataLoader(
val_dataset,
batch_size=batch_size,
shuffle=False,
num_workers=8,
pin_memory=True,
)
"""
Model, optimizer, loss setup
"""
model_dir = clean_path(os.path.join(working_dir, "model"))
optim = optim_const(model)
LOGGER.info("model set to {}".format(model))
LOGGER.info("optimizer set to {}".format(optim))
LOGGER.info("loss set to {}".format(loss))
LOGGER.info("devices set to {}".format(devices))
"""
Manager and config setup
"""
manager = ScheduledModifierManager.from_yaml(config_path)
logs_dir = clean_path(os.path.join(working_dir, "logs"))
loggers = [TensorBoardLogger(logs_dir), PythonLogger()]
optim = ScheduledOptimizer(
optim, model, manager, steps_per_epoch=len(train_loader), loggers=loggers
)
"""
Training and testing
"""
model, device, device_ids = model_to_device(model, devices)
trainer = ModuleTrainer(model, device, loss, optim, loggers=loggers)
tester = ModuleTester(model, device, loss, loggers=loggers, log_steps=-1)
epoch = -1
tester.run_epoch(val_loader, epoch=epoch)
for epoch in range(manager.max_epochs):
LOGGER.info("starting training epoch {}".format(epoch))
train_res = trainer.run_epoch(train_loader, epoch)
LOGGER.info("finished training epoch {}: {}".format(epoch, train_res))
val_res = tester.run_epoch(val_loader, epoch)
LOGGER.info("finished validation epoch {}: {}".format(epoch, val_res))
exporter = ModuleExporter(model, model_dir)
exporter.export_pytorch(optim, epoch)
for data in val_loader:
exporter.export_onnx(data)
def pb_to_onnx(
inputs: List[Union[str, tf_compat.Tensor]],
outputs: List[Union[str, tf_compat.Tensor]],
pb_path: str,
onnx_path: str,
opset: int = default_onnx_opset(),
custom_op_handlers=None,
extra_opset=None,
shape_override: Dict[str, List] = None,
):
"""
Export an ONNX format for the graph from PB format.
Should not be called within an active graph or session.
:param inputs: the inputs the graph should be created for,
can be either a list of names or a list of tensors
:param outputs: the outputs the graph should be created for,
can be either a list of names or a list of tensors
:param pb_path: path to the existing PB file
:param onnx_path: path to the output ONNX file
:param opset: ONNX opset
:param custom_op_handlers: dictionary of custom op handlers
:param extra_opset: list of extra opset's
:param shape_override: new shape to override
"""
try:
from tf2onnx import constants, optimizer, utils
from tf2onnx.tfonnx import process_tf_graph, tf_optimize
except ModuleNotFoundError:
raise ModuleNotFoundError(
"tf2onnx must be installed on the system before using export_onnx"
)
try:
from tf2onnx import tf_loader as loader
except Exception:
from tf2onnx import loader
pb_path = clean_path(pb_path)
if not os.path.exists(pb_path):
raise FileNotFoundError(
("no pb file for the model found at {}").format(pb_path)
)
inputs = [inp if isinstance(inp, str) else inp.name for inp in inputs]
outputs = [out if isinstance(out, str) else out.name for out in outputs]
graph_def, inputs, outputs = loader.from_graphdef(pb_path, inputs, outputs)
graph_def = tf_optimize(inputs, outputs, graph_def, fold_constant=True)
with tf_compat.Graph().as_default() as tf_graph:
tf_compat.import_graph_def(graph_def, name="")
with tf_compat.Session(graph=tf_graph):
graph = process_tf_graph(
tf_graph,
continue_on_error=False,
target=",".join(constants.DEFAULT_TARGET),
opset=opset,
custom_op_handlers=custom_op_handlers,
extra_opset=extra_opset,
shape_override=shape_override,
input_names=inputs,
output_names=outputs,
)
onnx_graph = optimizer.optimize_graph(graph)
model_proto = onnx_graph.make_model("converted from {}".format(pb_path))
onnx_path = clean_path(onnx_path)
create_parent_dirs(onnx_path)
utils.save_protobuf(onnx_path, model_proto)
