[docs]classTfliteMicro:"""This class wraps the TF-Lite Micro C++ library This class allows for loading a .tflite model file into the TF-Lite Micro (TFLM) C++ library and running inference using either the TFLM reference kernels or hardware accelerated kernels. """_model_lock=threading.Lock()_wrapper=None_logger:logging.Logger=None_logged_errors:List[str]=[]_accelerators:Dict[str,TfliteMicroAccelerator]={}_accelerator_paths:List[str]=[]
[docs]@staticmethoddefgit_hash()->str:"""Return the GIT hash of the MLTK repo used to compile the wrapper library"""wrapper=TfliteMicro._load_wrapper()returnwrapper.git_hash()
[docs]@staticmethoddefapi_version()->int:"""Return the TFLM API version number. This is used to ensure accelerator wrappers are compatible with this TFLM wrapper"""wrapper=TfliteMicro._load_wrapper()returnwrapper.api_version()
[docs]@staticmethoddefset_log_level(level:str)->str:"""Set the C++ wrapper logging level NOTE: This sets the level in the C++ wrapper, NOT the Python logger. Increasing the logging level can help with throughput as each log generated by the wrapper needs to be forwarded to the Python logger. Returns: The previous log level """wrapper=TfliteMicro._load_wrapper()prev_level=wrapper.get_log_level()ifnotwrapper.set_log_level(level):raiseRuntimeError(f'Failed to set MLTK log level to {level}')returnprev_level
[docs]@staticmethoddefget_log_level()->str:"""Return the C++ wrapper's logging level NOTE: This returns the C++ wrapper's logging level, NOT the Python logger. """wrapper=TfliteMicro._load_wrapper()returnwrapper.get_log_level()
[docs]@staticmethoddefset_logger(logger:logging.Logger):"""Set the wrapper's Python logger This logger will be invoked by the C++ wrapper's logging callback. """TfliteMicro._logger=logger
[docs]@staticmethoddefget_logger()->logging.Logger:"""Return the wrapper's Python logger"""# Just use the MLTK logger if no logger has been specifiedifTfliteMicro._loggerisNone:logger=get_mltk_logger()TfliteMicro._logger=loggerreturnTfliteMicro._logger
[docs]@staticmethoddefnormalize_accelerator_name(accelerator:str)->str:"""Given a case-insensitive accelerator name, normalize the name to the format used by the C++ library Returns: Normalized name of accelerator or None if accelerator is unknown """TfliteMicro._load_wrapper()ifacceleratorisNone:returnNonereturnget_case_insensitive(accelerator,TfliteMicro._accelerators)
[docs]@staticmethoddefget_supported_accelerators()->List[str]:"""Return a list of supported accelerators by name"""TfliteMicro._load_wrapper()return[xforxinTfliteMicro._accelerators]
[docs]@staticmethoddefaccelerator_is_supported(accelerator:str)->bool:"""Return if the given accelerator is supported"""TfliteMicro._load_wrapper()returnget_case_insensitive(accelerator,TfliteMicro._accelerators)isnotNone
[docs]@staticmethoddefload_tflite_model(model:Union[str,TfliteModel],accelerator:str=None,enable_profiler=False,enable_recorder=False,enable_tensor_recorder=False,force_buffer_overlap=False,runtime_buffer_size:int=None,**kwargs)->TfliteMicroModel:"""Load the TF-Lite Micro interpreter with the given .tflite model NOTE: - Only 1 model may be loaded at a time - You must call unload_model() when the model is no longer needed """wrapper=TfliteMicro._load_wrapper()ifacceleratorisnotNone:tflm_accelerator=TfliteMicro.get_accelerator(accelerator)ifhasattr(tflm_accelerator,'init_variant'):tflm_accelerator.init_variant()else:tflm_accelerator=NoneTfliteMicro._model_lock.acquire()try:tflite_model=_load_tflite_model(model)runtime_buffer_sizes=_retrieve_runtime_buffer_sizes(tflite_model,runtime_buffer_size=runtime_buffer_size,accelerator=accelerator)tflm_model=TfliteMicroModel(tflm_wrapper=wrapper,tflm_accelerator=tflm_accelerator,flatbuffer_data=tflite_model.flatbuffer_data,enable_profiler=enable_profiler,enable_recorder=enable_recorder,enable_tensor_recorder=enable_tensor_recorder,force_buffer_overlap=force_buffer_overlap,runtime_buffer_sizes=runtime_buffer_sizes,)except:# Release the model lock if an exception occurred while loading itTfliteMicro._model_lock.release()raisereturntflm_model
[docs]@staticmethoddefunload_model(model:TfliteMicroModel):"""Unload a previously loaded model"""accelerator=model.acceleratorifacceleratorisnotNone:ifhasattr(accelerator,'deinit_variant'):accelerator.deinit_variant()# pylint: disable=protected-accessifmodel._model_wrapper:model._model_wrapper.unload()delmodelTfliteMicro._model_lock.release()
[docs]@staticmethoddefprofile_model(model:Union[str,TfliteModel],accelerator:str=None,return_estimates=False,disable_simulator_backend=False,runtime_buffer_size=-1,# If runtime_buffer_size not given, determine the optimal memory sizeinput_data:Union[np.ndarray,List[np.ndarray]]=None,**kwargs):# -> ProfilingModelResults"""Profile the given model in the simulator and optionally determine metric estimates """frommltk.core.profiling_resultsimportProfilingModelResults,ProfilingLayerResulttflite_model=_load_tflite_model(model)tflm_model=TfliteMicro.load_tflite_model(model=tflite_model,accelerator=accelerator,enable_profiler=True,enable_recorder=True,runtime_buffer_size=runtime_buffer_size)try:renable_simulator_backend=Falsedisable_calculate_accelerator_cycles_only=Falsetflm_accelerator=tflm_model.acceleratorifdisable_simulator_backendand \
tflm_acceleratorisnotNoneand \
hasattr(tflm_accelerator,'set_simulator_backend_enabled'):renable_simulator_backend=Truetflm_accelerator.set_simulator_backend_enabled(False)ifhasattr(tflm_accelerator,'set_calculate_accelerator_cycles_only_enabled'):# For profiling, we only need the accelerator cycles# The simulator does not need to actually calculate valid output data# This greatly improves simulation latencydisable_calculate_accelerator_cycles_only=Truetflm_accelerator.set_calculate_accelerator_cycles_only_enabled(True)tflm_model_details=tflm_model.detailsifinput_dataisnotNone:ifisinstance(input_data,list):fori,vinenumerate(input_data):tflm_model.input(index=i,value=v)else:tflm_model.input(value=input_data)else:foriinrange(tflm_model.input_size):input_tensor=tflm_model.input(i)empty_tensor=np.zeros_like(input_tensor)tflm_model.input(i,value=empty_tensor)tflm_model.invoke()tflm_results=tflm_model.get_profiling_results()recorded_data=tflm_model.get_recorded_data()ifrenable_simulator_backend:tflm_accelerator.set_simulator_backend_enabled(True)ifdisable_calculate_accelerator_cycles_only:tflm_accelerator.set_calculate_accelerator_cycles_only_enabled(False)recorded_layers=recorded_data['layers']layer_results=[]forlayer_index,tflm_layer_resultinenumerate(tflm_results):tflite_layer=tflite_model.layers[layer_index]layer_err=tflm_model.get_layer_error(layer_index)layer_err_msg=Noneiflayer_errisNoneelselayer_err.msgdeltflm_layer_result['name']layer_result=ProfilingLayerResult(tflite_layer=tflite_layer,error_msg=layer_err_msg,**tflm_layer_result)layer_recorded_data=recorded_layers[layer_index]iflayer_index<len(recorded_layers)else{}updated=Truewhileupdated:updated=Falseforkey,valueinlayer_recorded_data.items():ifnotisinstance(value,(int,float,str)):tflite_layer.metadata[key]=valuelayer_recorded_data.pop(key)updated=Truebreaklayer_result.update(layer_recorded_data)layer_results.append(layer_result)finally:TfliteMicro.unload_model(tflm_model)model_details=tflm_model.details_add_memory_plan(tflite_model=tflite_model,recorded_data=recorded_data,model_details=model_details)results=ProfilingModelResults(model=tflite_model,accelerator=accelerator,runtime_memory_bytes=tflm_model_details.runtime_memory_size,layers=layer_results,model_details=model_details)# If we want to return estimates for metrics like:# CPU cycles and energyifreturn_estimates:# If accelerator=none# then just use the MVP accelerator's 'none' (i.e. CMSIS-only) estimatorsiftflm_acceleratorisNoneand'mvp'inTfliteMicro._accelerators:tflm_accelerator=TfliteMicro._accelerators['mvp']iftflm_acceleratorisnotNone:tflm_accelerator.estimate_profiling_results(results=results,**kwargs)returnresults
[docs]@staticmethoddefrecord_model(model:Union[str,TfliteModel],input_data:Union[np.ndarray,List[np.ndarray]]=None,accelerator:str=None,enable_accelerator_recorder=False,disable_simulator_backend=False,enable_tensor_recorder=True,return_model_details=False,update_input_model=False,layer_callback:Callable[[TfliteLayer],bool]=None,)->Union[List[TfliteLayer],Tuple[List[TfliteLayer],TfliteMicroModelDetails]]:"""Run one inference and record each model layer's input/output tensors Args: model: path to .tflite model file or TfliteModel instance input_data: Model input0 data as numpy array or list of numpy arrays for each model input accelerator: Optional accelerator to use for inference enable_accelerator_recorder: If enabled, record the data/instructions generated by the hardware accelerator The recorded data with be stored in each layers' metadata property, .e.g.: ``layer.metadata['accelerator_data']``. Each layers' recorded data is a dictionary with the entries specific to the hardware accelerator. disable_simulator_backend: Disable the simulator backend while running the accelerator recorder. This can greatly improve execution time, however, the generated data output (i.e. output tensors) is invalid enable_tensor_recorder: Record the input/output tensors of each layer return_model_details: Also return the recorded model's TfliteMicroModelDetails Return: If return_model_details=False, then return a list of TfliteLayers with the tensor data updated with the recorded values from the previous inference If return_model_details=True, then return a tuple(list(TfliteLayers), TfliteMicroModelDetails) """ifupdate_input_modelandnotisinstance(model,TfliteModel):raiseValueError('Input model must be a TfliteModel instance to use update_input_model=True')tflite_model=_load_tflite_model(model)tflm_model=TfliteMicro.load_tflite_model(model=tflite_model,accelerator=accelerator,enable_recorder=True,enable_tensor_recorder=enable_tensor_recorder,enable_profiler=False,runtime_buffer_size=16*1024*1024,# 16MB)disable_program_recorder=Falseifenable_accelerator_recorder:iftflm_model.acceleratorisNone:raiseValueError('Must provide accelerator when using enable_accelerator_recorder')tflm_model.accelerator.set_program_recorder_enabled(True)disable_program_recorder=Truereenable_simulator_backend=Falsereenable_simulator_accelerator_cycles=Falseifdisable_simulator_backend:iftflm_model.acceleratorisNone:raiseValueError('Must provide accelerator when using disable_simulator_backend')ifhasattr(tflm_model.accelerator,'set_simulator_backend_enabled'):reenable_simulator_backend=Truetflm_model.accelerator.set_simulator_backend_enabled(False)ifhasattr(tflm_model.accelerator,'set_calculate_accelerator_cycles_only_enabled'):tflm_model.accelerator.set_calculate_accelerator_cycles_only_enabled(True)iflayer_callback:tflm_model.set_layer_callback(functools.partial(_layer_callback_handler,tflite_model=tflite_model,callback=layer_callback))try:ifinput_dataisnotNone:ifisinstance(input_data,list):fori,vinenumerate(input_data):tflm_model.input(index=i,value=v)else:tflm_model.input(value=input_data)else:fori,inpinenumerate(tflite_model.inputs):d=np.zeros_like(inp.data)tflm_model.input(index=i,value=d)tflm_model.invoke()recorded_data=tflm_model.get_recorded_data()ifreenable_simulator_backend:tflm_model.accelerator.set_simulator_backend_enabled(True)ifreenable_simulator_accelerator_cycles:tflm_model.accelerator.set_calculate_accelerator_cycles_only_enabled(False)ifdisable_program_recorder:tflm_model.accelerator.set_program_recorder_enabled(False)retval=[]forlayer_index,recorded_layer_datainenumerate(recorded_data.get('layers',[])):# pylint: disable=protected-accesstf_layer=tflite_model.layers[layer_index]ifupdate_input_model \
elsecopy.deepcopy(tflite_model.layers[layer_index])retval.append(tf_layer)layer_err=tflm_model.get_layer_error(layer_index)tf_layer.metadata['error_msg']=Noneiflayer_errisNoneelselayer_err.msgforinput_index,input_bytesinenumerate(recorded_layer_data.get('inputs',[])):ifinput_index>=tf_layer.n_inputs:breakinput_tensor=tf_layer.inputs[input_index]ifinput_tensorisNone:continueinput_buf=np.frombuffer(input_bytes,dtype=input_tensor.dtype)ifinput_tensor.shape.flat_size>0:tf_layer.inputs[input_index]._data=np.reshape(input_buf,newshape=input_tensor.shape)else:tf_layer.inputs[input_index]._data=input_bufforoutput_index,output_bytesinenumerate(recorded_layer_data.get('outputs',[])):output_tensor=tf_layer.outputs[output_index]output_buf=np.frombuffer(output_bytes,dtype=output_tensor.dtype)ifoutput_tensor.shape.flat_size>0:tf_layer.outputs[output_index]._data=np.reshape(output_buf,newshape=output_tensor.shape)else:tf_layer.outputs[output_index]._data=output_bufforkey,valueinrecorded_layer_data.items():ifkeynotin('inputs','outputs'):tf_layer.metadata[key]=valueifreturn_model_details:model_details=tflm_model.details_add_memory_plan(tflite_model=tflite_model,recorded_data=recorded_data,model_details=model_details)finally:TfliteMicro.unload_model(tflm_model)ifreturn_model_details:returnretval,model_detailsreturnretval
[docs]@staticmethoddefadd_accelerator_path(path:str):"""Add an accelerator search path"""TfliteMicro._accelerator_paths.append(path)
[docs]@staticmethoddefregister_accelerator(accelerator:TfliteMicroAccelerator):"""Register a TFLM accelerator instance"""try:acc_api_version=accelerator.api_versionexceptExceptionase:# pylint:disable=raise-missing-fromraiseRuntimeError(f'Failed to load accelerator: {accelerator.name}, '+ \
f'failed to retrieve api version from wrapper, err: {e}')tflm_api_version=TfliteMicro.api_version()iftflm_api_version!=acc_api_version:raiseRuntimeError(f'Accelerator: {accelerator.name} not compatible, '+ \
f'accelerator API version ({acc_api_version}) != TFLM wrapper version ({tflm_api_version})')forvariantinaccelerator.variants:ifTfliteMicro.accelerator_is_supported(variant):raiseRuntimeError(f'Accelerator "{variant}" has already been registered')acc=copy.deepcopy(accelerator)acc.active_variant=variantTfliteMicro._accelerators[variant]=acc
[docs]@staticmethoddefget_accelerator(name:str)->TfliteMicroAccelerator:"""Return an instance to the specified accelerator wrapper"""TfliteMicro._load_wrapper()norm_accelerator=TfliteMicro.normalize_accelerator_name(name)ifnorm_acceleratorisNone:raiseValueError(f'Unknown accelerator: {name}. Known accelerators are: {", ".join(TfliteMicro.get_supported_accelerators())}')returnTfliteMicro._accelerators[norm_accelerator]
@staticmethoddef_load_wrapper():"""Load the TFLM C++ wrapper and return a refernce to the loaded module"""ifTfliteMicro._wrapperisnotNone:returnTfliteMicro._wrapper# Add this wrapper directory to the env PATH# This way, the wrapper DLL can find additional DLLs as necessarywrapper_dir=os.path.dirname(os.path.abspath(__file__))os.environ['PATH']=wrapper_dir+os.pathsep+os.environ['PATH']ifhasattr(os,'add_dll_directory'):os.add_dll_directory(wrapper_dir)# Import the TFLM C++ python wrapper# For more details, see:# <mltk root>/cpp/tflite_micro_wrappertry:TfliteMicro._wrapper=importlib.import_module('mltk.core.tflite_micro._tflite_micro_wrapper')except(ImportError,ModuleNotFoundError)ase:append_exception_msg(e,f'Failed to import the tflite_micro_wrapper C++ shared library.\n' \
'If you built the MLTK from source then this could mean you need to re-build the mltk package (e.g. "pip install -e .").\n' \
'If you\'re running from a pre-built MLTK package (e.g. "pip install silabs-mltk"),\n' \
f'ensure that the _tflite_micro_wrapper file exists at {wrapper_dir}.\n' \
'If the file does not exist, try installing, e.g.: pip install silabs-mltk --force-reinstall\n\n')raise# Initialize the wrapperTfliteMicro._wrapper.init()# Set the callback that will be invoked by the C++ library# log messagesTfliteMicro._wrapper.set_logger_callback(TfliteMicro._wrapper_logger_callback)TfliteMicro._load_accelerators()returnTfliteMicro._wrapper@staticmethoddef_load_accelerators():"""Load all the TFLM accelerators found in the search paths"""curdir=os.path.dirname(os.path.abspath(__file__))search_paths=[]search_paths.extend(TfliteMicro._accelerator_paths)search_paths.extend(as_list(get_user_setting('accelerator_paths')))search_paths.append(f'{curdir}/accelerators/mvp')# Check if any "<accelerator name>_mltk_accelerator.pth" files are found in the Python Libs directorypython_libs_dir=sysconfig.get_path('purelib')ifos.path.exists(python_libs_dir):forfninos.listdir(python_libs_dir):ifnotfn.endswith('_mltk_accelerator.pth'):continuepth_path=f'{python_libs_dir}/{fn}'withopen(pth_path,'r')asf:accelerator_package_base_dir=f.readline().strip()accelerator_name=fn[:-len('_mltk_accelerator.pth')]accelerator_dir=f'{accelerator_package_base_dir}/{accelerator_name}'# If the file does exist,# then add its path to the accelerator search pathifos.path.exists(accelerator_dir):search_paths.append(accelerator_dir)elifos.path.exists(f'{accelerator_dir}_wrapper'):search_paths.append(f'{accelerator_dir}_wrapper')forsearch_pathinsearch_paths:search_path=fullpath(search_path)init_py_path=f'{search_path}/__init__.py'ifnotos.path.exists(init_py_path):continueTfliteMicro._load_accelerator(search_path)TfliteMicro.register_accelerator(PlaceholderTfliteMicroAccelerator('cmsis'))@staticmethoddef_load_accelerator(accelerator_dir:str)->bool:"""Attempt to load an accelerator Python module in the given directory"""logger=TfliteMicro.get_logger()try:accelerator_module=import_module_at_path(accelerator_dir)exceptExceptionase:logger.debug(f'Failed to import {accelerator_dir}, err: {e}',exc_info=e)returnFalsetflm_accelerator=Noneforkeyindir(accelerator_module):value=getattr(accelerator_module,key)ifinspect.isclass(value)andissubclass(value,TfliteMicroAccelerator):# Create an accelerator instancetry:tflm_accelerator=value()breakexceptExceptionase:logger.warning(f'Accelerator module: {accelerator_dir} failed to initialize, err: \n{e}')returnFalseiftflm_acceleratorisNone:logger.debug(f'Accelerator module: {accelerator_dir} does not contain a TfliteMicroAccelerator class definition')returnFalsetry:TfliteMicro.register_accelerator(tflm_accelerator)exceptExceptionase:logger.warning(f'Failed to register accelerator: {accelerator_dir}, err: {e}')returnFalsereturnTrue@staticmethoddef_clear_logged_errors():"""Clear errors generated by C++ wrapper. This is used internally by the wrapper"""TfliteMicro._load_wrapper()TfliteMicro._logged_errors.clear()@staticmethoddef_get_logged_errors()->List[str]:"""Return errors generated by C++ wrapper as a list. This is used internally by the wrapper"""TfliteMicro._load_wrapper()returnTfliteMicro._logged_errors@staticmethoddef_get_logged_errors_str()->str:"""Return errors generated by C++ wrapper as a string. This is used internally by the wrapper"""return"\n".join(TfliteMicro._get_logged_errors())@staticmethoddef_wrapper_logger_callback(msg:str):""" This callback will be invoked by the TFLM C++ wrapper when it internally issues a log msg"""l=TfliteMicro.get_logger()iflisNone:returnerrs=TfliteMicro._logged_errorslevel=msg[:2].strip()msg=msg[2:].strip()iflevel=='D':l.debug(msg)eliflevel=='I':l.info(msg)eliflevel=='W':l.warning(msg)errs.append(msg)eliflevel=='E':l.error(msg)errs.append(msg)
def_load_tflite_model(model:Union[str,TfliteModel])->TfliteModel:ifisinstance(model,TfliteModel):returnmodelelifisinstance(model,str):ifnotmodel.endswith('.tflite')ornotos.path.exists(model):raiseValueError('Provided model must be a path to an existing .tflite file')returnTfliteModel.load_flatbuffer_file(model)else:raiseRuntimeError('Must provide TfliteModel or path to .tflite file')def_add_memory_plan(tflite_model:TfliteModel,recorded_data:List,model_details:TfliteMicroModelDetails):from.tflite_micro_memory_planimportTfliteMicroMemoryPlanrecorded_data_layers=recorded_data.get('layers',[])model_details._memory_plan=TfliteMicroMemoryPlan.create(# pylint: disable=protected-accessmemory_plan=recorded_data.get('memory_plan',[]),tflite_model=tflite_model,total_persistent_runtime_size=recorded_data.get('total_persistent_runtime_size',0),temp_runtime_sizes=list(x.get('temp_memory_used',0)forxinrecorded_data_layers),persistent_runtime_sizes=list(x.get('persistent_memory_used',0)forxinrecorded_data_layers))def_layer_callback_handler(tflite_model:TfliteModel,callback:Callable[[TfliteLayer],bool],index:int,outputs:List[bytes])->bool:tf_layer=copy.deepcopy(tflite_model.layers[index])foroutput_index,output_bytesinenumerate(outputs):output_tensor=tf_layer.outputs[output_index]output_buf=np.frombuffer(output_bytes,dtype=output_tensor.dtype)ifoutput_tensor.shape.flat_size>0:tf_layer.outputs[output_index]._data=np.reshape(output_buf,newshape=output_tensor.shape)# pylint: disable=protected-accesselse:tf_layer.outputs[output_index]._data=output_buf# pylint: disable=protected-accessreturncallback(tf_layer)def_retrieve_runtime_buffer_sizes(tflite_model:TfliteModel,accelerator:str,runtime_buffer_size:int,)->List[int]:frommltk.core.tflite_model_parametersimportTfliteModelParametersruntime_buffer_sizes=[0]try:memory_spec=TfliteModelParameters.load_from_tflite_model(tflite_model,tag=f'{accelerator}_memory_spec')except:memory_spec={}ifmemory_spec:runtime_buffer_sizes=memory_spec.get('sizes',[0])ifruntime_buffer_size:runtime_buffer_sizes[0]=runtime_buffer_sizereturnruntime_buffer_sizes
Important: We use cookies only for functional and traffic analytics.
We DO NOT use cookies for any marketing purposes. By using our site you acknowledge you have read and understood our Cookie Policy.