def import_all_params(require_success: bool = True) -> None:
"""Imports all ModelParams to add them to the global registry.
Because the BUILD rule may selectively depend on a subset of task params,
we use a try-except to guard around every import.
Args:
require_success: Whether we require the underlying lingvo ImportAllParams()
function to succeed or not.
"""
success = lingvo_model_imports.ImportAllParams(_TASK_ROOT,
_TASK_DIRS,
require_success=False)
if require_success and not success:
raise ValueError(
'Could not import any task params. Make sure task params '
'are linked into the binary.')
for i, k in enumerate(fetch_keys)
if k in self._fetches.keys()])
for k in kwargs:
if k not in self._feeds:
raise KeyError(
"%s is not in the list of available feeds. Available keys: %s" %
(k, list(self._feeds.keys())))
feeds = {self._feeds[k]: v for k, v in six.iteritems(kwargs)}
run_options = tf.RunOptions(report_tensor_allocations_upon_oom=False)
if session_run_options:
run_options = session_run_options
fetched_results = self._RunWithValidSession(
tf.Session.run, valid_fetches, feed_dict=feeds, options=run_options)
results = [None] * len(fetch_keys)
for i, fetch in zip(valid_fetch_idxs, fetched_results):
results[i] = fetch
return results
def main(_):
import IPython # pylint: disable=g-import-not-at-top
IPython.start_ipython(argv=["--colors", "NoColor"], user_ns=globals())
if __name__ == "__main__":
model_imports.ImportAllParams()
tf.app.run(main)
def GetAllRegisteredClasses():
model_imports.ImportAllParams()
return _ModelRegistryHelper.GetAllRegisteredClasses()
