"""
Manages all the data and build processes associated with the TensorFlow graph.
The TensorFlow graph is the symbolic description of the computations in the
network, which will be executed by the simulator.
"""
from collections import OrderedDict, defaultdict
import functools
import inspect
import itertools
import logging
import warnings
from nengo import Connection, Process, Ensemble
from nengo.builder.operator import TimeUpdate, SimPyFunc, Reset
from nengo.builder.processes import SimProcess
from nengo.config import ConfigError
from nengo.ensemble import Neurons
from nengo.exceptions import SimulationError, ValidationError
from nengo.neurons import Direct
from nengo.utils.magic import decorator
import numpy as np
import tensorflow as tf
from nengo_dl import builder, graph_optimizer, signals, utils, tensor_node, config
from nengo_dl.compat import (
tf_compat,
SparseMatrix,
is_sparse,
make_process_state,
make_process_step,
)
logger = logging.getLogger(__name__)
[docs]@decorator
def with_self(wrapped, instance, args, kwargs):
"""A decorator that can be used to ensure that any ops created within the
wrapped method will be added to the TensorGraph object's graph."""
with instance.graph.as_default(), instance.graph.device(instance.device):
return wrapped(*args, **kwargs)
[docs]class TensorGraph:
"""
Manages the construction of the TensorFlow symbolic computation graph.
Parameters
----------
model : `~nengo.builder.Model`
Pre-built Nengo model describing the network to be simulated
dt : float
Length of a simulator timestep, in seconds
unroll_simulation : int
Unroll simulation loop by explicitly building ``unroll_simulation``
iterations into the computation graph
dtype : ``tf.DType``
Floating point precision to use for simulation
minibatch_size : int
The number of simultaneous inputs that will be passed through the
network
device : None or ``"/cpu:0"`` or ``"/gpu:[0-n]"``
Device on which to execute computations (if None then uses the
default device as determined by TensorFlow)
progress : `.utils.ProgressBar`
Progress bar for optimization stage
"""
def __init__(
self, model, dt, unroll_simulation, dtype, minibatch_size, device, progress
):
self.model = model
self.dt = dt
self.unroll = unroll_simulation
self.dtype = dtype
self.minibatch_size = minibatch_size
self.device = device
self.graph = tf.Graph()
self.signals = signals.SignalDict(self.dtype, self.minibatch_size)
self.inference_only = config.get_setting(model, "inference_only", False)
# find invariant inputs (nodes that don't receive any input other
# than the simulation time). we'll compute these outside the simulation
# and feed in the result.
if self.model.toplevel is None:
self.invariant_inputs = OrderedDict()
else:
self.invariant_inputs = OrderedDict(
(n, n.output)
for n in self.model.toplevel.all_nodes
if n.size_in == 0 and not isinstance(n, tensor_node.TensorNode)
)
# filter unused operators
# remove TimeUpdate because it is executed as part of the simulation
# loop, not part of the step plan. remove input nodes because they
# are executed outside the simulation.
node_processes = [
n.output for n in self.invariant_inputs if isinstance(n.output, Process)
]
operators = [
op
for op in self.model.operators
if not (
isinstance(op, TimeUpdate)
or (isinstance(op, SimPyFunc) and op.x is None)
or (
isinstance(op, SimProcess)
and op.input is None
and op.process in node_processes
)
)
]
# mark trainable signals
self.mark_signals()
logger.info("Initial plan length: %d", len(operators))
# apply graph simplification functions
simplifications = config.get_setting(
model,
"simplifications",
[
graph_optimizer.remove_constant_copies,
graph_optimizer.remove_unmodified_resets,
graph_optimizer.remove_zero_incs,
graph_optimizer.remove_identity_muls,
],
)
with progress.sub("operator simplificaton", max_value=None):
old_operators = []
while len(old_operators) != len(operators) or any(
x is not y for x, y in zip(operators, old_operators)
):
old_operators = operators
for simp in simplifications:
operators = simp(operators)
# group mergeable operators
planner = config.get_setting(model, "planner", graph_optimizer.tree_planner)
with progress.sub("merging operators", max_value=None):
plan = planner(operators)
# TODO: we could also merge operators sequentially (e.g., combine
# a copy and dotinc into one op), as long as the intermediate signal
# is only written to by one op and read by one op
# order signals/operators to promote contiguous reads
sorter = config.get_setting(model, "sorter", graph_optimizer.order_signals)
with progress.sub("ordering signals", max_value=None):
sigs, self.plan = sorter(plan, n_passes=10)
# create base arrays and map Signals to TensorSignals (views on those
# base arrays)
with progress.sub("creating signals", max_value=None):
self.create_signals(sigs)
logger.info("Optimized plan length: %d", len(self.plan))
logger.info("Number of base arrays: %d", len(self.base_arrays_init))
# initialize op builder
build_config = builder.BuildConfig(
inference_only=self.inference_only,
lif_smoothing=config.get_setting(self.model, "lif_smoothing"),
cpu_only=self.device == "/cpu:0" or not utils.tf_gpu_installed,
)
self.op_builder = builder.Builder(
self.plan, self.graph, self.signals, build_config
)
[docs] @with_self
def build(self, progress):
"""
Constructs a new graph to simulate the model.
progress : `.utils.ProgressBar`
Progress bar for construction stage
"""
self.target_phs = {}
self.outputs = {}
self.optimizers = {}
# create these constants once here for reuse in different operators
self.signals.dt = tf.constant(self.dt, self.dtype)
self.signals.dt_val = self.dt # store the actual value as well
self.signals.zero = tf.constant(0, self.dtype)
self.signals.one = tf.constant(1, self.dtype)
if not self.inference_only:
# this variable controls behaviour in the simulation that is
# conditional on whether we are doing training or inference
self.signals.training = tf_compat.placeholder(
tf.bool, shape=(), name="training"
)
# variable to track training step
self.training_step = tf_compat.train.get_or_create_global_step()
else:
self.training_step = None
# create base arrays
sub = progress.sub("creating base arrays")
self.base_vars = OrderedDict()
unique_ids = defaultdict(int)
for k, (v, trainable) in sub(self.base_arrays_init.items()):
name = "%s_%s_%s_%d" % (
v.dtype,
"_".join(str(x) for x in v.shape),
trainable,
unique_ids[(v.dtype, v.shape, trainable)],
)
unique_ids[(v.dtype, v.shape, trainable)] += 1
# we initialize all the variables from placeholders, and then
# feed in the initial values when the init op is called. this
# prevents TensorFlow from storing large constants in the graph
# def, which can cause problems for large models
ph = tf_compat.placeholder(v.dtype, v.shape, name="%s_init" % name)
if trainable:
with tf_compat.variable_scope("trainable_vars", reuse=False):
var = tf_compat.get_variable(
name, initializer=ph, trainable=True, use_resource=False
)
else:
with tf_compat.variable_scope("local_vars", reuse=False):
var = tf_compat.get_local_variable(
name, initializer=ph, trainable=False, use_resource=False
)
self.base_vars[k] = (var, ph, v)
logger.debug("created base arrays")
logger.debug([str(x[0]) for x in self.base_vars.values()])
# set up invariant inputs
sub = progress.sub("building inputs")
self.build_inputs(sub)
# pre-build stage
with progress.sub("pre-build stage", max_value=len(self.plan)) as sub:
self.op_builder.pre_build(sub)
# build stage
with progress.sub("build stage", max_value=len(self.plan) * self.unroll) as sub:
self.build_loop(sub)
# ops for initializing variables (will be called by simulator)
trainable_vars = tf_compat.trainable_variables()
if not self.inference_only:
trainable_vars.append(self.training_step)
self.trainable_init_op = tf_compat.variables_initializer(trainable_vars)
self.local_init_op = tf_compat.local_variables_initializer()
self.global_init_op = tf_compat.variables_initializer(
[v for v in tf_compat.global_variables() if v not in trainable_vars]
)
self.constant_init_op = tf_compat.variables_initializer(
tf_compat.get_collection("constants")
)
# logging
logger.info(
"Number of reads: %d", sum(x for x in self.signals.read_types.values())
)
for x in self.signals.read_types.items():
logger.info(" %s: %d", *x)
logger.info(
"Number of writes: %d", sum(x for x in self.signals.write_types.values())
)
for x in self.signals.write_types.items():
logger.info(" %s: %d", *x)
[docs] def build_step(self, progress):
"""
Build the operators that execute a single simulation timestep
into the graph.
Parameters
----------
progress : `.utils.ProgressBar`
Progress bar for loop construction
Returns
-------
probe_tensors : list of ``tf.Tensor``
The Tensor objects representing the data required for each model
Probe
side_effects : list of ``tf.Tensor``
The output Tensors of computations that may have side-effects
(e.g., `~nengo.Node` functions), meaning that they
must be executed each time step even if their output doesn't appear
to be used in the simulation
"""
# manually build TimeUpdate. we don't include this in the plan,
# because loop variables (`step`) are (semi?) pinned to the CPU, which
# causes the whole variable to get pinned to the CPU if we include
# `step` as part of the normal planning process.
self.signals.time = tf.cast(self.signals.step, self.dtype) * self.signals.dt
# build operators
side_effects = self.op_builder.build(progress)
logger.debug("collecting probe tensors")
probe_tensors = []
for p in self.model.probes:
probe_sig = self.model.sig[p]["in"]
if probe_sig in self.signals:
# TODO: better solution to avoid the forced_copy
# we need to make sure that probe reads occur before the
# probe value is overwritten on the next timestep. however,
# just blocking on the sliced value (probe_tensor) doesn't
# work, because slices of variables don't perform a
# copy, so the slice can be "executed" and then the value
# overwritten before the tensorarray write occurs. what we
# really want to do is block until the probe_arrays.write
# happens, but you can't block on probe_arrays (and blocking on
# probe_array.flow doesn't work, although I think it should).
# so by adding the copy here and then blocking on the copy, we
# make sure that the probe value is read before it can be
# overwritten.
probe_tensors.append(
self.signals.gather(self.signals[probe_sig], force_copy=True)
)
else:
# if a probe signal isn't in sig_map, that means that it isn't
# involved in any simulator ops. so we know its value never
# changes, and we'll just return a constant containing the
# initial value.
if probe_sig.minibatched:
init_val = np.tile(
probe_sig.initial_value[..., None], (1, self.minibatch_size)
)
else:
init_val = probe_sig.initial_value
probe_tensors.append(tf.constant(init_val, dtype=self.dtype))
logger.debug("=" * 30)
logger.debug("build_step complete")
logger.debug("probe_tensors %s", [str(x) for x in probe_tensors])
logger.debug("side_effects %s", [str(x) for x in side_effects])
return probe_tensors, side_effects
[docs] def build_loop(self, progress):
"""
Build simulation loop.
Parameters
----------
progress : `.utils.ProgressBar`
Progress bar for loop construction
"""
def loop_condition(step, stop, *_):
return step < stop
def loop_body(step, stop, loop_i, probe_arrays, base_vars):
# fill in signals.bases (note: we need to do this here because we
# need to use the versions of the base vars from inside the
# loop, not the static variables in self.base_vars)
assert len(self.signals.bases) == 0
for i, key in enumerate(
itertools.chain(
self.base_vars.keys(), self.signals.internal_vars.keys()
)
):
self.signals.bases[key] = base_vars[i]
for iter in range(self.unroll):
logger.debug("BUILDING ITERATION %d", iter)
with self.graph.name_scope("iteration_%d" % iter):
# note: nengo step counter is incremented at the beginning
# of the timestep
step += 1
self.signals.step = step
# fill in invariant input data
for n in self.input_ph:
self.signals.scatter(
self.signals[self.model.sig[n]["out"]],
self.input_ph[n][loop_i],
)
# build the operators for a single step
# note: we tie things to the `loop_i` variable so that we
# can be sure the other things we're tying to the
# simulation step (side effects and probes) from the
# previous timestep are executed before the next step
# starts
# note2: we use the variable scope to make sure that we
# aren't accidentally creating new variables for
# unrolled iterations (this is really only a concern
# with TensorNodes)
scope = tf_compat.variable_scope(
tf_compat.get_variable_scope(), reuse=iter > 0
)
with self.graph.control_dependencies([loop_i]), scope:
probe_tensors, side_effects = self.build_step(progress)
# copy probe data to array
for i, p in enumerate(probe_tensors):
if config.get_setting(
self.model,
"keep_history",
default=True,
obj=self.model.probes[i],
):
probe_arrays[i] = probe_arrays[i].write(loop_i, p)
else:
probe_arrays[i] = tf.cond(
pred=tf.equal(step, stop),
true_fn=lambda p=p: probe_arrays[i].write(0, p),
false_fn=lambda: probe_arrays[i],
)
# need to make sure that any operators that could have side
# effects run each timestep, so we tie them to the loop
# increment. we also need to make sure that all the probe
# reads happen before those values get overwritten on the
# next timestep
with self.graph.control_dependencies(side_effects + probe_tensors):
loop_i += 1
base_vars = tuple(self.signals.bases.values())
return step, stop, loop_i, probe_arrays, base_vars
self.step_var = tf_compat.placeholder(tf.int32, shape=(), name="step")
self.stop_var = tf_compat.placeholder(tf.int32, shape=(), name="stop")
loop_i = tf.constant(0)
probe_arrays = [
tf.TensorArray(self.dtype, clear_after_read=True, size=0, dynamic_size=True)
for _ in self.model.probes
]
# build simulation loop
loop_vars = (
self.step_var,
self.stop_var,
loop_i,
probe_arrays,
tuple(
x[0]._ref() if isinstance(x[0], tf.Variable) else x[0]
for x in self.base_vars.values()
)
+ tuple(x._ref() for x in self.signals.internal_vars.values()),
)
loop_vars = tf.while_loop(
cond=loop_condition,
body=loop_body,
loop_vars=loop_vars,
parallel_iterations=1,
back_prop=not self.inference_only,
)
self.steps_run = loop_vars[2]
self.probe_arrays = OrderedDict()
for p, a in zip(self.model.probes, loop_vars[3]):
x = a.stack()
if self.model.sig[p]["in"].minibatched:
x = tf.transpose(a=x, perm=np.roll(np.arange(x.get_shape().ndims), 1))
else:
x = tf.expand_dims(x, 0)
self.probe_arrays[p] = x
[docs] def build_optimizer_func(self, optimizer, objective):
"""
Adds elements into the graph to execute the given optimizer.
Parameters
----------
optimizer : ``tf.train.Optimizer``
Instance of a TensorFlow optimizer class
objective : dict of {`~nengo.Probe`: callable or ``None``}
The objective to be minimized. This is a dictionary mapping Probes
to functions
``f(output, target) -> loss`` that consume the actual output and
target output for the given probe(s) and return a ``tf.Tensor``
representing a scalar loss value. The function may also accept a
single argument ``f(output) -> loss`` if targets are not required.
Some common objective functions can be found in
`nengo_dl.objectives`.
Passing ``None`` as the probe value (instead of a callable)
indicates that the error is being computed outside the simulation,
and the value passed for that probe in ``data`` directly specifies
the output error gradient.
If multiple probes are specified as the key, then the corresponding
output/target values will be passed as a list to the objective
function.
The overall loss value being minimized will be the sum across all
the objectives specified.
Returns
-------
apply_optimizer : callable
A function that builds the operators required to implement the
given optimizer update. Generally this function will then be
passed to `~.build_outputs`.
Notes
-----
This function caches its outputs, so if it is called again with the
same arguments then it will return the previous function. This avoids
building duplicates of the same operations over and over. This can
also be important functionally, e.g. if the optimizer has internal
state like momentum. By caching the output we ensure that subsequent
calls share the same internal state.
"""
key = (optimizer, frozenset(objective.items()))
try:
# return the cached optimizer function if it exists
return self.optimizers[key]
except KeyError:
pass
# note: the standard workflow is that sim.train calls
# build_optimizer_func to get this function. it then passes the
# function to run_batch, which calls build_outputs to actually
# build these operations into the graph. we do this somewhat
# indirect method so that everything passes through build_output,
# allowing us to consolidate certain logic there (like capturing
# new variables)
def apply_optimizer(outputs, targets):
# note: we don't actually use outputs/targets, because the same
# data is pulled implicitly from `objective` below.
# we just check that outputs and targets match up with
# objective, to make sure there's nothing weird going on.
if not isinstance(outputs, tuple):
outputs = (outputs,)
if not isinstance(targets, tuple):
targets = (targets,)
assert set(outputs) == set(self.probe_arrays[p] for p in objective)
assert set(targets) == set(self.target_phs[p] for p in objective)
agg_method = tf.AggregationMethod.DEFAULT
grads = []
vars = tf_compat.trainable_variables()
# compute loss
# note: we drop the `None` items in objective, because we
# want to treat those as direct gradients (rather than
# returning the probe value, which is the standard behaviour for
# build_outputs)
loss, _ = self.build_outputs(
{k: v for k, v in objective.items() if v is not None}
)
# compute gradients wrt loss
if len(loss) > 0:
# reduce loss to a scalar
loss = tf.reduce_sum(
input_tensor=[tf.reduce_sum(input_tensor=v) for v in loss.values()]
)
grads.append(
tf.gradients(ys=loss, xs=vars, aggregation_method=agg_method)
)
# add in any gradients where the user directly specified the output
# error grad
for p, g in objective.items():
if g is None:
grads.append(
tf.gradients(
ys=self.probe_arrays[p],
xs=vars,
grad_ys=self.target_phs[p],
aggregation_method=agg_method,
)
)
# combine gradients for each variable
if len(grads) == 1:
grads = grads[0]
else:
grads = [tf.reduce_sum(input_tensor=gs, axis=0) for gs in zip(*grads)]
opt_op = optimizer.apply_gradients(
zip(grads, tf_compat.trainable_variables())
)
with tf.control_dependencies([opt_op]):
new_step = tf_compat.assign_add(
self.training_step, tf.constant(1, dtype=tf.int64)
)
return new_step, loss
self.optimizers[key] = apply_optimizer
return apply_optimizer
[docs] @with_self
def build_outputs(self, outputs):
"""
Adds elements into the graph to compute the given outputs.
Parameters
----------
outputs : dict of {(tuple of) `~nengo.Probe`: callable or None}
The output function to be applied to each probe or group of probes.
The function can accept one argument (the output of that probe) or
two (output and target values for that probe). If a tuple of
Probes are given as the key, then those output/target parameters
will be the corresponding tuple of probe/target values. The
function should return a ``tf.Tensor`` or tuple of Tensors
representing the output we want from those probes. If ``None`` is
given instead of a function then the output will simply be the
output value from the corresponding probes.
Returns
-------
output_vals : dict of {(tuple of) `~nengo.Probe`: \
(tuple of) ``tf.Tensor``}
Tensors representing the result of applying the output functions
to the probes.
new_vars_init : ``tf.Tensor`` or None
Initialization op for any new variables created when building
the outputs.
Notes
-----
This function caches its outputs, so if it is called again with the
same arguments then it will return the previous Tensors. This avoids
building duplicates of the same operations over and over. This can
also be important functionally, e.g. if the outputs have internal
state. By caching the output we ensure that subsequent
calls share the same internal state.
"""
key = frozenset(outputs.items())
try:
# return the cached outputs if they exist
return self.outputs[key], None
except KeyError:
pass
output_vals = {}
new_vars = []
for probes, out in outputs.items():
is_tuple = isinstance(probes, tuple)
probe_arrays = (
tuple(self.probe_arrays[p] for p in probes)
if is_tuple
else self.probe_arrays[probes]
)
if out is None:
# return probe output value
output_vals[probes] = probe_arrays
elif callable(out):
# look up number of arguments for function
spec = inspect.getfullargspec(out)
nargs = len(spec.args)
# don't count keyword arguments
if spec.defaults is not None:
nargs -= len(spec.defaults)
# don't count self argument for methods or callable classes
out_func = out.func if isinstance(out, functools.partial) else out
if inspect.ismethod(out_func) or not inspect.isroutine(out_func):
nargs -= 1
# build function arguments
if nargs == 1:
args = [probe_arrays]
elif nargs == 2:
for p in probes if is_tuple else (probes,):
# create a placeholder for the target values if one
# hasn't been created yet
if p not in self.target_phs:
self.target_phs[p] = tf_compat.placeholder(
self.dtype,
(self.minibatch_size, None, p.size_in),
name="%s_ph" % utils.sanitize_name(p),
)
target_phs = (
tuple(self.target_phs[p] for p in probes)
if is_tuple
else self.target_phs[probes]
)
args = [probe_arrays, target_phs]
else:
raise ValidationError(
"Output functions must accept 1 or 2 arguments; '%s' "
"takes %s arguments"
% (utils.function_name(out, sanitize=False), nargs),
"outputs",
)
# apply output function
with tf_compat.variable_scope(utils.function_name(out)) as scope:
output_vals[probes] = out(*args)
# collect any new variables from building the outputs
for collection in [
tf_compat.GraphKeys.GLOBAL_VARIABLES,
tf_compat.GraphKeys.LOCAL_VARIABLES,
"gradient_vars",
]:
new_vars.extend(scope.get_collection(collection))
else:
raise ValidationError("Outputs must be callable or None)", "outputs")
new_vars_init = (
tf_compat.variables_initializer(new_vars) if len(new_vars) > 0 else None
)
self.outputs[key] = output_vals
return output_vals, new_vars_init
[docs] @with_self
def build_post(self, sess, rng):
"""
Executes post-build processes for operators (after the graph has
been constructed and session/variables initialized).
Note that unlike other build functions, this is called every time
the simulator is reset.
Parameters
----------
sess : ``tf.Session``
The TensorFlow session for the simulator
rng : `~numpy.random.mtrand.RandomState`
Seeded random number generator
"""
# build input functions (we need to do this here, because in the case
# of processes these functions depend on the rng, and need to be be
# rebuilt on reset)
self.input_funcs = {}
for n, output in self.invariant_inputs.items():
if isinstance(output, np.ndarray):
self.input_funcs[n] = output
elif isinstance(output, Process):
state = make_process_state(output, (n.size_in,), (n.size_out,), self.dt)
self.input_funcs[n] = [
make_process_step(
output,
(n.size_in,),
(n.size_out,),
self.dt,
output.get_rng(rng),
state,
)
for _ in range(self.minibatch_size)
]
elif n.size_out > 0:
self.input_funcs[n] = [
utils.align_func((n.size_out,), self.dtype)(output)
]
else:
# a node with no inputs and no outputs, but it can still
# have side effects
self.input_funcs[n] = [output]
# execute post_build on all the op builders
self.op_builder.post_build(sess, rng)
[docs] @with_self
def build_summaries(self, summaries):
"""
Adds ops to collect summary data for the given objects.
Parameters
----------
summaries : list of dict or \
`~nengo.Connection` or \
`~nengo.Ensemble` or \
`~nengo.ensemble.Neurons` or \
``tf.Tensor``}
List of objects for which we want to collect data. Object can be a
Connection (in which case data on weights will be collected),
Ensemble (encoders), Neurons (biases), a dict of
``{probe: objective}`` that indicates a loss function that will
be tracked, or a pre-built summary tensor.
Returns
-------
op : ``tf.Tensor``
Merged summary op for the given summaries
"""
summary_ops = []
inits = []
with tf.device("/cpu:0"):
for obj in summaries:
if isinstance(obj, dict):
# overall loss
loss, init = self.build_outputs(obj)
if init is not None:
inits.append(init)
summary_ops.append(
tf_compat.summary.scalar(
"loss",
tf.reduce_sum(
input_tensor=[
tf.reduce_sum(input_tensor=v) for v in loss.values()
]
),
family="loss",
)
)
if len(obj) > 1:
# get loss for each probe
for p, t in loss.items():
summary_ops.append(
tf_compat.summary.scalar(
utils.sanitize_name("Probe_%s_loss" % p.label),
tf.reduce_sum(input_tensor=t),
family="loss",
)
)
elif isinstance(obj, (Ensemble, Neurons, Connection)):
if isinstance(obj, Ensemble):
param = "encoders"
name = "Ensemble_%s" % obj.label
elif isinstance(obj, Neurons):
param = "bias"
name = "Ensemble.neurons_%s" % obj.ensemble.label
elif isinstance(obj, Connection):
param = "weights"
name = "Connection_%s" % obj.label
summary_ops.append(
tf_compat.summary.histogram(
utils.sanitize_name("%s_%s" % (name, param)),
self.get_tensor(self.model.sig[obj][param]),
)
)
elif isinstance(obj, tf.Tensor):
# we assume that obj is a summary op
summary_ops.append(obj)
else:
raise SimulationError("Unknown summary object: %s" % obj)
return (
tf_compat.summary.merge(summary_ops),
(None if len(inits) == 0 else inits),
)
[docs] @with_self
def get_tensor(self, sig):
"""
Returns a Tensor corresponding to the given Signal.
Parameters
----------
sig : `~nengo.builder.Signal`
A signal in the model
Returns
-------
tensor : ``tf.Tensor``
Tensor containing the value of the given Signal
"""
tensor_sig = self.signals[sig]
base = self.base_vars[tensor_sig.key][0]
if "while/" in tensor_sig.tf_indices.name:
# rebuild tf indices outside the while loop
tensor_sig._tf_indices = None
return tf.gather(base, tensor_sig.tf_indices)
[docs] def mark_signals(self):
"""
Mark all the signals in ``self.model`` according to whether they
represent trainable parameters of the model (parameters that can be
optimized by deep learning methods).
Trainable parameters include connection weights, ensemble encoders, and
neuron biases. Unless one of those signals is targeted by a Nengo
learning rule (otherwise the learning rule update conflicts with the
deep learning optimization).
Users can manually specify whether signals are trainable or not using
the config system (e.g.,
``net.config[nengo.Ensemble].trainable = False``)
"""
def get_trainable(net_config, obj, network_trainable):
"""Looks up the current value of ``obj.trainable``."""
if self.inference_only:
return False
try:
if obj in net_config.params:
# priority #1: instance config
trainable = net_config[obj].trainable
elif network_trainable is not 1: # noqa: F632
# priority #2: network setting
trainable = network_trainable
else:
# priority #3: class config
trainable = net_config[obj].trainable
except (ConfigError, AttributeError):
trainable = network_trainable
# we return 1 if trainable isn't configured, since the default is
# for everything to be trainable but we want to be able to
# distinguish whether something was specifically set to be
# trainable (True) or just defaulting to trainable (1)
return 1 if trainable is None else trainable
def mark_network(net_config, net, network_trainable):
"""Recursively marks the signals for objects within each
subnetwork."""
for subnet in net.networks:
mark_network(
net_config,
subnet,
get_trainable(net_config, subnet, network_trainable),
)
# encoders and biases are trainable
for ens in net.ensembles:
ens_trainable = get_trainable(net_config, ens, network_trainable)
self.model.sig[ens]["encoders"].trainable = ens_trainable
self.model.sig[ens]["encoders"].minibatched = False
if not isinstance(ens.neuron_type, Direct):
neurons_trainable = get_trainable(
net_config, ens.neurons, network_trainable
)
if neurons_trainable is 1: # noqa: F632
neurons_trainable = ens_trainable
self.model.sig[ens.neurons]["bias"].trainable = neurons_trainable
self.model.sig[ens.neurons]["bias"].minibatched = False
# connection weights are trainable
for conn in net.connections:
# note: this doesn't include probe connections, since they
# aren't added to the network
self.model.sig[conn]["weights"].trainable = get_trainable(
net_config, conn, network_trainable
)
self.model.sig[conn]["weights"].minibatched = False
# parameters can't be modified by an online Nengo learning rule
# and offline training at the same time. (it is possible in
# theory, but it complicates things a lot and is probably not a
# common use case). we also make those signals minibatched
# (they wouldn't be normally), because we want to be able to
# learn independently in each minibatch
for conn in net.connections:
rule = conn.learning_rule
if rule is not None:
if isinstance(rule, dict):
rule = list(rule.values())
elif not isinstance(rule, list):
rule = [rule]
for r in rule:
if r.modifies in ("weights", "decoders"):
obj = conn
attr = "weights"
elif r.modifies == "encoders":
obj = conn.post_obj
attr = "encoders"
else:
raise NotImplementedError
if self.model.sig[obj][attr].trainable is True:
warnings.warn(
"%s has a learning rule and is also set "
"to be trainable; this is likely to "
"produce strange training behaviour." % obj
)
else:
self.model.sig[obj][attr].trainable = False
self.model.sig[obj][attr].minibatched = True
if self.model.toplevel is None:
warnings.warn(
"No top-level network in model; assuming no trainable parameters",
UserWarning,
)
else:
net_config = self.model.toplevel.config
mark_network(
net_config,
self.model.toplevel,
get_trainable(net_config, self.model.toplevel, 1),
)
# the connections to connection probes are not trainable, but
# also not minibatched
probe_seeds = [self.model.seeds[p] for p in self.model.probes]
for obj, seed in self.model.seeds.items():
if isinstance(obj, Connection) and seed in probe_seeds:
self.model.sig[obj]["weights"].trainable = False
self.model.sig[obj]["weights"].minibatched = False
# fill in defaults for all other signals
# signals are not trainable by default, and views take on the
# properties of their bases
for op in self.model.operators:
for sig in op.all_signals:
if not hasattr(sig.base, "trainable"):
sig.base.trainable = False
if not hasattr(sig.base, "minibatched"):
sig.base.minibatched = not sig.base.trainable
if not hasattr(sig, "trainable"):
sig.trainable = sig.base.trainable
if not hasattr(sig, "minibatched"):
sig.minibatched = sig.base.minibatched
[docs] def create_signals(self, sigs):
"""
Groups signal data together into larger arrays, and represent each
individual signal as a slice into that array.
Parameters
----------
sigs : list of `~nengo.builder.Signal`
Base signals arranged into the order in which they should reside in
memory (e.g., output from `.graph_optimizer.order_signals`)
"""
float_type = self.dtype.as_numpy_dtype
base_arrays = OrderedDict()
curr_keys = {}
sig_idxs = {s: i for i, s in enumerate(sigs)}
# find the non-overlapping partitions of the signals
breaks = []
diff = defaultdict(int)
for ops in self.plan:
# note: we don't include Resets, otherwise the big reset block
# overrides most of the partitioning
if not isinstance(ops[0], Reset):
for i in range(len(ops[0].all_signals)):
op_sigs = [op.all_signals[i].base for op in ops]
idxs = [sig_idxs[s] for s in op_sigs]
diff[op_sigs[np.argmin(idxs)]] += 1
diff[op_sigs[np.argmax(idxs)]] -= 1
# find the partition points in signal list
open = 0
for i, s in enumerate(sigs):
if s in diff:
open += diff[s]
if open == 0:
breaks += [i + 1]
logging.debug("partitions")
logging.debug(
"\n%s", "".join("|" if i in breaks else " " for i in range(len(sigs)))
)
# create all the base signals
for i, sig in enumerate(sigs):
assert sig not in self.signals
assert not sig.is_view
if i in breaks:
# start a new array for all current bases
for k in curr_keys:
curr_keys[k] = object()
# convert to appropriate dtype
if np.issubdtype(sig.dtype, np.floating):
dtype = float_type
elif np.issubdtype(sig.dtype, np.integer):
dtype = np.int32
elif np.issubdtype(sig.dtype, np.bool_):
dtype = sig.dtype
else:
raise NotImplementedError("Unsupported signal dtype")
if is_sparse(sig):
# for sparse tensors, what we care about is the shape of the
# underlying data, not the full matrix
shape = (sig.initial_value.size,)
else:
# resize scalars to length 1 vectors
shape = sig.shape if sig.shape != () else (1,)
# parameters of signal that affect the base array
array_params = (dtype, shape[1:], sig.trainable, sig.minibatched)
# key used to map signals to base arrays
if array_params not in curr_keys:
curr_keys[array_params] = object()
key = curr_keys[array_params]
initial_value = sig.initial_value
if is_sparse(sig):
if isinstance(initial_value, SparseMatrix):
initial_value = initial_value.data
else:
initial_value = initial_value.tocoo().data
initial_value = initial_value.astype(dtype, copy=False)
# broadcast scalars up to full size
if initial_value.shape == ():
initial_value = np.resize(initial_value, shape)
if sig.minibatched:
# duplicate along minibatch dimension
initial_value = np.tile(
initial_value[..., None],
tuple(1 for _ in shape) + (self.minibatch_size,),
)
if key in base_arrays:
base_arrays[key][0].append(initial_value)
base_arrays[key][2] += shape[0]
else:
base_arrays[key] = [[initial_value], sig.trainable, shape[0]]
n = base_arrays[key][-1]
indices = np.arange(n - shape[0], n)
tensor_sig = self.signals.get_tensor_signal(
indices, key, dtype, shape, sig.minibatched, label=sig.name, signal=sig
)
logger.debug("created base signal")
logger.debug(sig)
logger.debug(tensor_sig)
for key in base_arrays:
arrs, t, _ = base_arrays[key]
base_arrays[key] = (np.concatenate(arrs, axis=0), t)
# add any signal views to the sig_map
all_views = [
sig
for ops in self.plan
for op in ops
for sig in op.all_signals
if sig.is_view
]
for sig in all_views:
if sig.size == sig.base.size:
# reshape view
self.signals[sig] = self.signals[sig.base].reshape(sig.shape)
else:
if sig.shape[1:] != sig.base.shape[1:]:
# TODO: support this?
raise NotImplementedError("Slicing on axes > 0 is not supported")
# slice view
assert np.all([x == 1 for x in sig.elemstrides[1:]])
start = sig.elemoffset
stride = sig.elemstrides[0]
stop = start + sig.size * stride
if stop < 0:
stop = None
self.signals[sig] = self.signals[sig.base][slice(start, stop, stride)]
# error checking
for sig, tensor_sig in self.signals.items():
# tensorsignal shapes should match signal shapes
assert (
tensor_sig.shape == (sig.size,)
if is_sparse(sig)
else (sig.shape if sig.shape != () else (1,))
)
# tensorsignal values should match signal values
initial_value = sig.initial_value
if is_sparse(sig):
if isinstance(initial_value, SparseMatrix):
initial_value = initial_value.data
else:
initial_value = initial_value.tocoo().data
if sig.minibatched:
initial_value = initial_value[..., None]
assert np.allclose(
base_arrays[tensor_sig.key][0][tensor_sig.indices],
initial_value.astype(dtype),
)
logger.debug("base arrays")
logger.debug(
"\n".join(
[
str((k, v.dtype, v.shape, trainable))
for k, (v, trainable) in base_arrays.items()
]
)
)
self.base_arrays_init = base_arrays
