# Copyright (c) 2020 Sony Corporation. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import os
from collections import OrderedDict
import nnabla as nn
from nnabla.utils.save import save
from .parameter import Parameter
from hydra import utils
[docs]
class Module(object):
r"""Module base for all nnabla neural network modules.
Your models should also subclass this class. Modules can also contain
other Modules, allowing to nest them in a tree structure.
"""
def __init__(self, name=''):
self._name = name
if os.environ.get('NNABLA_NAS_MIXEDOP_FAST_MODE') is not None:
self._call_create = self.call
self.call = self._call_cached
self._train_output = None
self._infer_output = None
@property
def name(self):
r"""
The name of the module.
Returns:
string: the name of the module
"""
return self._name
@property
def modules(self):
r"""Return an `OrderedDict` containing immediate modules."""
if '_modules' not in self.__dict__:
self.__dict__['_modules'] = OrderedDict()
return self._modules
@property
def parameters(self):
r"""Return an `OrderedDict` containing immediate parameters."""
if '_parameters' not in self.__dict__:
self.__dict__['_parameters'] = OrderedDict()
return self._parameters
@property
def training(self):
r"""The training mode of module."""
if '_training' not in self.__dict__:
self.__dict__['_training'] = True
return self._training
@training.setter
def training(self, mode):
self.__dict__['_training'] = mode
for _, m in self.modules.items():
m.training = mode
@property
def need_grad(self):
r"""Whether the module needs gradient."""
if '_need_grad' not in self.__dict__:
self.__dict__['_need_grad'] = True
return self._need_grad
@need_grad.setter
def need_grad(self, mode):
self.__dict__['_need_grad'] = mode
for _, m in self.modules.items():
m.need_grad = mode
@property
def is_active(self):
r"""Whether the module was called."""
if '_is_active' not in self.__dict__:
self.__dict__['_is_active'] = True
return self._is_active
@is_active.setter
def is_active(self, mode):
self.__dict__['_is_active'] = mode
@property
def input_shapes(self):
r"""Return a list of input shapes used during `call` function."""
if '_input_shapes' not in self.__dict__:
self.__dict__['_input_shapes'] = list()
return self._input_shapes
@input_shapes.setter
def input_shapes(self, v):
setattr(self, '_input_shapes', v)
def _get_need_grad_state(self):
from nnabla import parameter
# TODO: Ideally want to have get_current_no_grad_state() in parameter module?
no_grad = parameter.current_no_grad
if no_grad:
return False
return self.need_grad
def __getattr__(self, name):
if name in self.modules:
return self.modules[name]
if name in self.parameters:
need_grad_state = self._get_need_grad_state()
p = self.parameters[name]
if not need_grad_state and p.need_grad:
return p.get_unlinked_variable(need_grad=need_grad_state)
return p
return object.__getattr__(self, name)
def __setattr__(self, name, value):
def remove_from(*dicts):
for d in dicts:
d.pop(name, None)
remove_from(self.__dict__, self.modules, self.parameters)
if isinstance(value, Module):
self.modules[name] = value
elif isinstance(value, Parameter):
self.parameters[name] = value
else:
object.__setattr__(self, name, value)
def __delattr__(self, name):
if name in self.parameters:
del self.parameters[name]
elif name in self.modules:
del self.modules[name]
else:
object.__delattr__(self, name)
[docs]
def apply(self, memo=None, **kargs):
r"""Helper for setting property recursively, then returns self."""
if memo is None:
memo = set()
if self not in memo:
memo.add(self)
for key, value in kargs.items():
setattr(self, key, value)
for module in self.modules.values():
module.apply(memo, **kargs)
return self
[docs]
def get_modules(self, prefix='', memo=None):
r"""Returns an iterator over all modules in the network, yielding
both the name of the module as well as the module itself.
Args:
prefix (str, optional): Additional prefix to name modules.
Defaults to ''.
memo (dict, optional): Memorize all parsed modules.
Defaults to None.
Yields:
(str, Module): a submodule.
"""
if memo is None:
memo = set()
if self not in memo:
memo.add(self)
yield prefix, self
for name, module in self.modules.items():
submodule_prefix = prefix + ('/' if prefix else '') + name
for m in module.get_modules(submodule_prefix, memo):
yield m
[docs]
def get_parameters(self, grad_only=False):
r"""Return an `OrderedDict` containing all parameters in the module.
Args:
grad_only (bool, optional): If need_grad=True is required.
Defaults to False.
Returns:
OrderedDict: A dictionary containing parameters of module.
"""
params = OrderedDict()
for prefix, module in self.get_modules():
if grad_only and not module.need_grad:
continue
for name, p in module.parameters.items():
if grad_only and not p.need_grad:
continue
key = prefix + ('/' if prefix else '') + name
params[key] = p
return params
[docs]
def set_parameters(self, params, raise_if_missing=False):
r"""Set parameters for the module.
Args:
params (OrderedDict): The parameters which will be loaded.
raise_if_missing (bool, optional): Raise exception if some
parameters are missing. Defaults to `False`.
Raises:
ValueError: Parameters are not found.
"""
for prefix, module in self.get_modules():
for name, p in module.parameters.items():
key = prefix + ('/' if prefix else '') + name
if key in params:
p.d = params[key].d.copy()
nn.logger.info(f'`{key}` loaded.')
elif raise_if_missing:
raise ValueError(
f'A child module {name} cannot be found in '
'{this}. This error is raised because '
'`raise_if_missing` is specified '
'as True. Please turn off if you allow it.')
[docs]
def save_parameters(self, path, params=None, grad_only=False):
r"""Saves the parameters to a file.
Args:
path (str): Absolute path to file.
params (OrderedDict, optional): An `OrderedDict` containing
parameters. If params is `None`, then the current parameters
will be saved.
grad_only (bool, optional): If need_grad=True is required for
parameters which will be saved. Defaults to False.
"""
params = params or self.get_parameters(grad_only)
nn.save_parameters(path, params)
[docs]
def load_parameters(self, path, raise_if_missing=False):
r"""Loads parameters from a file with the specified format.
Args:
path (str): Relative path to the parameter file (based on the original working directory).
raise_if_missing (bool, optional): Raise exception if some
parameters are missing. Defaults to `False`.
"""
with nn.parameter_scope('', OrderedDict()):
# adjust path because hydra changes the working directory
load_path = utils.to_absolute_path(path)
nn.load_parameters(load_path)
params = nn.get_parameters(grad_only=False)
self.set_parameters(params, raise_if_missing=raise_if_missing)
@property
def modules_to_profile(self):
r"""Returns a list with the modules that will be profiled when the
Profiler/Estimator functions are called. All other modules in the
network will not be profiled.
"""
raise NotImplementedError
[docs]
def get_latency(self, estimator, active_only=True):
"""
Function to use to calc latency
This function needs to work based on the graph
Parameters:
estimator: a graph-based estimator
active_only: get latency of active modules only
Returns:
latencies: list of all latencies of each module
accum_lat: total sum of latencies of all modules
"""
accum_lat = 0
latencies = {}
for mi in self.get_net_modules(active_only=active_only):
if type(mi) in self.modules_to_profile:
inp = [nn.Variable((1,) + si[1:]) for si in mi.input_shapes]
out = mi.call(*inp)
latencies[mi.name] = estimator.predict(out)
accum_lat += latencies[mi.name]
return latencies, accum_lat
[docs]
def get_latency_by_mod(self, estimator, active_only=True):
"""
*** Note: This function is deprecated. Use get_latency() ***
Function to use to calc latency
This function needs to work based on the module
Parameters:
estimator: a module-based estimator
active_only: get latency of active modules only
Returns:
latencies: list of all latencies of each module
accum_lat: total sum of latencies of all modules
"""
accum_lat = 0
latencies = {}
for mi in self.get_net_modules(active_only=active_only):
if type(mi) in self.modules_to_profile:
latencies[mi.name] = estimator.predict(mi)
accum_lat += latencies[mi.name]
return latencies, accum_lat
[docs]
def save_net_nnp(self, path, inp, out, calc_latency=False,
func_real_latency=None, func_accum_latency=None,
save_params=None):
"""
Saves whole net as one nnp Calc whole net (real) latency (using
e.g.Nnabla's [Profiler]) Calculate also layer-based latency The modules
are discovered using the nnabla graph of the whole net The latency is
then calculated based on each individual module's nnabla graph (e.g.
[LatencyGraphEstimator])
Args:
path: absolute path
inp: input of the created network
out: output of the created network
calc_latency: flag for calc latency
func_real_latency: function to use to calc actual latency
func_accum_latency: function to use to calc accum. latency,
this is, dissecting the network layer by layer
using the graph of the network, calculate the
latency for each layer and add up all these results.
"""
batch_size = inp.shape[0]
name = self.name if (hasattr(self, 'name')
and self.name) else 'results'
filename = os.path.join(path, name + '.nnp')
os.makedirs(os.path.dirname(filename), exist_ok=True)
name_for_nnp = self.name if (
hasattr(self, 'name') and self.name) else 'empty'
contents = {'networks': [{'name': name_for_nnp,
'batch_size': batch_size,
'outputs': {"y'": out},
'names': {'x': inp}}],
'executors': [{'name': 'runtime',
'network': name_for_nnp,
'data': ['x'],
'output': ["y'"]}]}
if save_params and 'no_image_normalization' in save_params:
contents['executors'][0]['no_image_normalization'] = save_params['no_image_normalization']
save(filename, contents, variable_batch_size=False)
if calc_latency:
acc_latency = func_accum_latency.get_estimation(out)
filename = path + name + '.acclat'
with open(filename, 'w') as f:
print(acc_latency.__str__(), file=f)
func_real_latency.run()
real_latency = float(func_real_latency.result['forward_all'])
filename = path + name + '.realat'
with open(filename, 'w') as f:
print(real_latency.__str__(), file=f)
return real_latency, acc_latency
else:
return 0.0, 0.0
[docs]
def save_modules_nnp(self, path, active_only=False,
calc_latency=False,
func_latency=None
):
"""
Saves all modules of the network as individual nnp files, using folder
structure given by name convention. The modules are extracted going
over the module list, not over the graph structure. The latency is
then calculated based on each individual module's nnabla graph (e.g.
[LatencyGraphEstimator])
Args:
path
active_only: if True, only active modules are saved
calc_latency: flag for calc latency
func_latency: function to use to calc latency of
each of the extracted modules
This function needs to work based on the graph
"""
accum_lat = 0.0
mods = self.get_net_modules(active_only=active_only)
for mi in mods:
if type(mi) in self.modules_to_profile:
if len(mi.input_shapes) == 0:
continue
pass
inp = [nn.Variable((1,) + si[1:]) for si in mi.input_shapes]
out = mi.call(*inp)
filename = path + mi.name + '.nnp'
pathname = os.path.dirname(filename)
upper_pathname = os.path.dirname(pathname)
if not os.path.exists(upper_pathname):
os.mkdir(upper_pathname)
if not os.path.exists(pathname):
os.mkdir(pathname)
d_dict = {str(i): inpi for i, inpi in enumerate(inp)}
d_keys = [str(i) for i, inpi in enumerate(inp)]
name_for_nnp = mi.name if (mi.name != '') else 'empty'
contents = {'networks': [{'name': name_for_nnp,
'batch_size': 1,
'outputs': {'out': out},
'names': d_dict}],
'executors': [{'name': 'runtime',
'network': name_for_nnp,
'data': d_keys,
'output': ['out']}]}
if hasattr(mi, '_scope_name'):
with nn.parameter_scope(mi._scope_name):
save(filename, contents, variable_batch_size=False)
else:
save(filename, contents, variable_batch_size=False)
if calc_latency:
latency = func_latency.get_estimation(out)
filename = path + mi.name + '.acclat'
with open(filename, 'w') as f:
print(latency.__str__(), file=f)
accum_lat += latency
return accum_lat
[docs]
def save_modules_nnp_by_mod(self, path, active_only=False,
calc_latency=False,
func_latency=None,
):
"""
*** Note: This function is deprecated. Use save_modules_nnp() ***
Saves all modules of the network as individual nnp files,
using folder structure given by name convention.
The modules are extracted going over the module list, not
over the graph structure.
The latency is then calculated using the module themselves
(e.g. [LatencyEstimator])
Args:
path
active_only: if True, only active modules are saved
calc_latency: flag for calc latency
func_latency: function to use to calc latency of
each of the extracted modules
This function needs to work based on the modules
"""
accum_lat = 0.0
mods = self.get_net_modules(active_only=active_only)
for mi in mods:
if type(mi) in self.modules_to_profile:
if len(mi.input_shapes) == 0:
continue
pass
inp = [nn.Variable((1,) + si[1:]) for si in mi.input_shapes]
out = mi.call(*inp)
filename = path + mi.name + '.nnp'
pathname = os.path.dirname(filename)
upper_pathname = os.path.dirname(pathname)
if not os.path.exists(upper_pathname):
os.mkdir(upper_pathname)
if not os.path.exists(pathname):
os.mkdir(pathname)
d_dict = {str(i): inpi for i, inpi in enumerate(inp)}
d_keys = [str(i) for i, inpi in enumerate(inp)]
name_for_nnp = mi.name if (mi.name != '') else 'empty'
contents = {'networks': [{'name': name_for_nnp,
'batch_size': 1,
'outputs': {'out': out},
'names': d_dict}],
'executors': [{'name': 'runtime',
'network': name_for_nnp,
'data': d_keys,
'output': ['out']}]}
if hasattr(mi, '_scope_name'):
with nn.parameter_scope(mi._scope_name):
save(filename, contents, variable_batch_size=False)
else:
save(filename, contents, variable_batch_size=False)
if calc_latency:
latency = func_latency.get_estimation(mi)
filename = path + mi.name + '.acclat'
with open(filename, 'w') as f:
print(latency.__str__(), file=f)
accum_lat += latency
return accum_lat
[docs]
def calc_latency_all_modules(self, path, graph, func_latency=None):
"""
Calculate the latency for each of the modules in a graph.
The modules are extracted using the graph structure information.
The latency is then calculated based on each individual module's
nnabla graph.
It also saves the accumulated latency of all modules.
Args:
path
graph:
func_latency: function to use to calc latency of
each of the modules
This function needs to work based on the graph
"""
import nnabla.function as Function
from nnabla_nas.utils.estimator.latency import Profiler
from nnabla.context import get_current_context
from nnabla.logger import logger
func_latency._visitor.reset()
graph.visit(func_latency._visitor)
total_latency = 0.0
idx = 0
for func in func_latency._visitor._functions:
args = [func.info.type_name] + \
[str(inp.shape) for inp in func.inputs] + \
[str(func.info.args)]
key = '-'.join(args)
ff = getattr(Function, func.info.type_name)(get_current_context(),
**func.info.args)
if key not in func_latency.memo:
try: # run profiler
nnabla_vars = [nn.Variable(inp.shape,
need_grad=inp.need_grad)
for inp in func.inputs]
runner = Profiler(
ff(*nnabla_vars),
device_id=func_latency._device_id,
ext_name=func_latency._ext_name,
n_run=func_latency._n_run,
outlier=func_latency._outlier,
max_measure_execution_time=func_latency._max_measure_execution_time, # noqa: E501
time_scale=func_latency._time_scale,
n_warmup=func_latency._n_warmup
)
runner.run()
latency = float(runner.result['forward_all'])
except Exception as err:
latency = 0.0
logger.warning(f'Latency calculation failed: {key}')
logger.warning(str(err))
func_latency.memo[key] = latency
else:
latency = func_latency.memo[key]
total_latency += latency
# save latency of this layer (name: id_XXX_{key}.acclat)
filename = path + '/id_' + str(idx) + '_' + key + '.acclat'
pathname = os.path.dirname(filename)
upper_pathname = os.path.dirname(pathname)
if not os.path.exists(upper_pathname):
os.mkdir(upper_pathname)
if not os.path.exists(pathname):
os.mkdir(pathname)
idx += 1
with open(filename, 'w') as f:
print(latency.__str__(), file=f)
# save accum latency of all layers
filename = path + '.acclat'
with open(filename, 'w') as f:
print(total_latency.__str__(), file=f)
return total_latency
[docs]
def convert_npp_to_onnx(self, path, opset='opset_11'):
"""
Finds all nnp files in the given path and its subfolders and converts
them to ONNX For this to run smoothly, nnabla_cli must be installed and
added to your python path.
Args:
path
opset
The actual bash shell command used is::
> find <DIR> -name '*.nnp' -exec echo echo {} \|
awk -F \\. \'\{print \"nnabla_cli convert -b 1 -d opset_11 \"\$0\" \"\$1\"\.\"\$2\"\.onnx\"\}\' \; | sh | sh
which, for each file found with find, outputs the following::
> echo <FILE>.nnp | awk -F \. '{print "nnabla_cli convert -b 1 -d opset_11 "$0" "$1"."$2".onnx"}' # noqa: E501,W605
which, for each file, generates the final conversion command::
> nnabla_cli convert -b 1 -d opset_11 <FILE>.nnp <FILE>.nnp.onnx
"""
os.system('find ' + path + ' -name "*.nnp" -exec echo echo {} \|' # noqa: E501,W605
' awk -F \\. \\\'{print \\\"nnabla_cli convert -b 1 -d ' + opset + # noqa: E501,W605
' \\\"\$0\\\" \\\"\$1\\\"\.\\\"\$2\\\"\.onnx\\\"}\\\' \; | sh | sh' # noqa: E501,W605
)
def __str__(self):
r"""Return str representtation of the module."""
main_str = f'{self.__class__.__name__}(' + self.extra_repr()
sub_str = ''
for key, module in self.modules.items():
m_repr = str(module).split('\n')
head = [self.extra_format().format(key) + ': ' + m_repr.pop(0)]
tail = [m_repr.pop()] if len(m_repr) else []
m_repr = [' ' * 2 + line for line in (head + m_repr + tail)]
sub_str += '\n' + '\n'.join(m_repr)
main_str += sub_str + ('\n' if sub_str else '') + ')'
return main_str
def __call__(self, *args, **kwargs):
self.input_shapes = [x.shape for x in args]
return self.call(*args, **kwargs)
[docs]
def call(self, *args, **kwargs):
r"""Implement the call of module. Inputs should only be Variables."""
raise NotImplementedError
def _call_cached(self, *args, **kwargs):
if self.training:
if self._train_output is None:
self._train_output = self._call_create(*args, **kwargs)
return self._train_output
else:
if self._infer_output is None:
self._infer_output = self._call_create(*args, **kwargs)
return self._infer_output