#########################################################
# Copyright (C) 2022 SiMa Technologies, Inc.
#
# This material is SiMa proprietary and confidential.
#
# This material may not be copied or distributed without
# the express prior written permission of SiMa.
#
# All rights reserved.
#########################################################
# Code owner: Joey Chou
#########################################################
"""Entry points for writing AwesomeNet to files and reading it from files."""
import yaml
import numpy as np
import os
import zipfile
import json
from io import BytesIO, StringIO
from afe.common_utils import get_afe_git_commit_id
from afe.ir.serializer import serializer
from afe.ir.serializer.visualizer import create_netron_graph
from afe.ir.net import AwesomeNet
from sima_utils.logging import sima_logger
[docs]
def save_awesomenet(net: AwesomeNet, model_name: str, output_directory: str = "") -> None:
"""
Saves an AwesomeNet to a yaml and npz file. Overwrites existing files.
:param net: AwesomeNet to save.
:param model_name: The name of the model, used in file names.
:param output_directory: The directory to save the files to.
"""
# get output file paths
sima_file_name = model_name if model_name.endswith('.sima') else model_name + '.sima'
json_file_name = sima_file_name + '.json'
yaml_file = sima_file_name + ".yaml"
numpy_file = sima_file_name + ".npz"
netron_file = os.path.join(output_directory, json_file_name)
zip_path = os.path.join(output_directory, sima_file_name)
# encode data
yaml_data, numpy_data = serializer.encode_awesomenet(net)
commit_id = get_afe_git_commit_id()
yaml_data = serializer.wrap_yaml_document(yaml_data, "Network", "1", commit_id)
netron_dict = create_netron_graph(net, model_name)
# Encode and save the dictionary in netron format to a file
encoder = serializer.AwesomeEncoder()
with open(netron_file, "w") as write_file:
json.dump(netron_dict, write_file, default=encoder.default, indent=4)
# Save numpy data and yaml data to their buffers
numpy_io = BytesIO()
np.savez(numpy_io, **numpy_data)
yaml_io = StringIO()
yaml.safe_dump(yaml_data, yaml_io)
# Create a zip archive and write down the buffers
try:
archive_context_manager = zipfile.ZipFile(zip_path, mode="w")
except (OSError, zipfile.BadZipFile) as e:
raise sima_logger.UserFacingException("Unable to open file '{}' for writing".format(zip_path)) from e
with archive_context_manager as archive:
archive.writestr(yaml_file, yaml_io.getvalue())
archive.writestr(numpy_file, numpy_io.getvalue())
sima_logger.sima_log_info(f"Saved model: {zip_path}")
[docs]
def load_awesomenet(model_name: str, network_directory: str = "") -> AwesomeNet:
"""
Loads an AwesomeNet from yaml and npz files
:param model_name: The name of the model, used in file names.
:param network_directory: The directory to load the files from.
"""
# get file paths
sima_file_name = model_name if model_name.endswith('.sima') else model_name + '.sima'
yaml_file = sima_file_name + ".yaml"
numpy_file = sima_file_name + ".npz"
zip_path = os.path.join(network_directory, sima_file_name)
# Open the zip archive and load yaml data and numpy data stored in byte format.
try:
archive_context_manager = zipfile.ZipFile(zip_path, mode="r")
except (OSError, zipfile.BadZipFile) as e:
raise sima_logger.UserFacingException("Unable to read file '{}'".format(zip_path)) from e
with archive_context_manager as archive:
with archive.open(yaml_file) as y_file:
yaml_data = y_file.read()
yaml_bytes = BytesIO(yaml_data)
yaml_data = yaml.safe_load(yaml_bytes)
with archive.open(numpy_file) as n_file:
numpy_data = n_file.read()
bytes_numpy = BytesIO(numpy_data)
numpy_data = np.load(bytes_numpy, allow_pickle=False)
sima_logger.sima_log_info(f"Loading model from file: {zip_path}")
# decode data
commit_id = get_afe_git_commit_id()
yaml_data = serializer.unwrap_yaml_document(yaml_data, "Network", "1", commit_id)
return serializer.decode_awesomenet(yaml_data, numpy_data)