Source code for nengo_dl.config

from nengo import Connection, Ensemble, Network, ensemble
from nengo.builder import Model
from nengo.exceptions import ConfigError, NetworkContextError
from nengo.params import BoolParam, Parameter


def configure_settings(**kwargs):
    """
    Pass settings to ``nengo_dl`` by setting them as parameters on the
    top-level Network config.

    The settings are passed as keyword arguments to ``configure_settings``;
    e.g., to set ``trainable`` use ``configure_settings(trainable=True)``.

    Parameters
    ----------
    trainable : bool or None
        Adds a parameter to Nengo Ensembles/Connections/Networks that controls
        whether or not they will be optimized by :meth:`.Simulator.train`.
        Passing ``None`` will use the default ``nengo_dl`` trainable settings,
        or True/False will override the default for all objects.  In either
        case trainability can be further configured on a per-object basis (e.g.
        ``net.config[my_ensemble].trainable = True``.  See `the documentation
        <https://www.nengo.ai/nengo-dl/training.html#choosing-which-elements-to-optimize>`_
        for more details.
    planner : graph planning algorithm
        Pass one of the `graph planners
        <https://www.nengo.ai/nengo-dl/graph_optimizer.html>`_ to change the
        default planner.
    sorter : signal sorting algorithm
        Pass one of the `sort algorithms
        <https://www.nengo.ai/nengo-dl/graph_optimizer.html>`_ to change the
        default sorter.
    simplifications: list of graph simplification functions
        Pass a list of `graph simplification functions
        <https://www.nengo.ai/nengo-dl/graph_optimizer.html>`_ to change the
        default simplifications applied.
    session_config: dict
        Config options passed to ``tf.Session`` initialization (e.g., to change
        the `GPU memory allocation method
        <https://www.tensorflow.org/guide/using_gpu#allowing_gpu_memory_growth>`_
        pass ``{"gpu_options.allow_growth": True}``).
    inference_only : bool
        Set to True if the network will only be run in inference mode (i.e.,
        no calls to :meth:`.Simulator.train`).  This may result in a small
        increase in the inference speed.
    lif_smoothing : float
        If specified, use the smoothed :class:`~.neurons.SoftLIFRate` neuron
        model, with the given smoothing parameter (``sigma``),
        to compute the gradient for :class:`~nengo:nengo.LIF` neurons (as
        opposed to using :class:`~nengo:nengo.LIFRate`).
    dtype : ``tf.DType``
        Set the floating point precision for simulation values.
    """

    # get the toplevel network
    if len(Network.context) > 0:
        config = Network.context[0].config
    else:
        raise NetworkContextError(
            "`configure_settings` must be called within a Network context "
            "(`with nengo.Network(): ...`)")

    try:
        params = config[Network]
    except ConfigError:
        config.configures(Network)
        params = config[Network]

    for attr, val in kwargs.items():
        if attr == "trainable":
            for obj in (Ensemble, Connection, ensemble.Neurons, Network):
                try:
                    obj_params = config[obj]
                except ConfigError:
                    config.configures(obj)
                    obj_params = config[obj]

                obj_params.set_param("trainable", BoolParam("trainable", val,
                                                            optional=True))
        elif attr in ("planner", "sorter", "simplifications",
                      "session_config", "inference_only", "lif_smoothing",
                      "dtype"):
            params.set_param(attr, Parameter(attr, val))
        else:
            raise ConfigError("%s is not a valid config parameter" % attr)


def get_setting(model, setting, default=None):
    """
    Returns config settings (created by :func:`.configure_settings`).

    Parameters
    ----------
    model : :class:`~nengo:nengo.builder.Model` or \
            :class:`~nengo:nengo.Network`
        Built model or Network containing all the config settings.
    setting : str
        Name of the config option to return
    default
        The default value to return if config option not set

    Returns
    -------
    Value of ``setting`` if it has been specified, else ``default``.
    """

    if isinstance(model, Model):
        if model.toplevel is None:
            return default
        model = model.toplevel

    try:
        return getattr(model.config[model], setting, default)
    except ConfigError:
        return default