def __init__(self, ops, signals, config):
super().__init__(ops, signals, config)
self.J_data = signals.combine([op.J for op in ops])
self.output_data = signals.combine([op.output for op in ops])
self.state_data = [
signals.combine([op.states[i] for op in ops])
for i in range(len(ops[0].states))
]
self.prev_result = []
def neuron_step_math(dt, J, *states): # pragma: no cover (runs in TF)
output = None
J_offset = 0
state_offset = [0 for _ in states]
for op in ops:
# slice out the individual state vectors from the overall
# array
op_J = J[:, J_offset:J_offset + op.J.shape[0]]
J_offset += op.J.shape[0]
op_states = []
for j, s in enumerate(op.states):
op_states += [
states[j][:,
state_offset[j]:state_offset[j] + s.shape[0]]
]
state_offset[j] += s.shape[0]
# call step_math function
# note: `op_states` are views into `states`, which will
# be updated in-place
mini_out = []
for j in range(signals.minibatch_size):
# blank output variable
neuron_output = np.zeros(op.output.shape,
self.output_data.dtype)
op.neurons.step_math(dt, op_J[j], neuron_output,
*[s[j] for s in op_states])
mini_out += [neuron_output]
neuron_output = np.stack(mini_out, axis=0)
# concatenate outputs
if output is None:
output = neuron_output
else:
output = np.concatenate((output, neuron_output), axis=1)
return (output, ) + states
self.neuron_step_math = neuron_step_math
self.neuron_step_math.__name__ = utils.sanitize_name("_".join(
[repr(op.neurons) for op in ops]))
def build_pre(self, signals, config):
super().build_pre(signals, config)
self.time_data = signals[self.ops[0].t].reshape(())
self.input_data = (None if self.ops[0].input is None else
signals.combine([op.input for op in self.ops]))
self.output_data = signals.combine([op.output for op in self.ops])
self.state_data = [
signals.combine([list(op.state.values())[i] for op in self.ops])
for i in range(len(self.ops[0].state))
]
self.mode = "inc" if self.ops[0].mode == "inc" else "update"
self.prev_result = []
# `merged_func` calls the step function for each process and
# combines the result
def merged_func(time, *input_state): # pragma: no cover (runs in TF)
if not hasattr(self, "step_fs"):
raise SimulationError("build_post has not been called for %s" %
self)
if self.input_data is None:
input = None
state = input_state
else:
input = input_state[0]
state = input_state[1:]
# update state in-place (this will update the state values
# inside step_fs)
for i, s in enumerate(state):
self.step_states[i][...] = s
input_offset = 0
func_output = []
for i, op in enumerate(self.ops):
if op.input is not None:
input_shape = op.input.shape[0]
func_input = input[:,
input_offset:input_offset + input_shape]
input_offset += input_shape
mini_out = []
for j in range(signals.minibatch_size):
x = [] if op.input is None else [func_input[j]]
mini_out += [self.step_fs[i][j](*([time] + x))]
func_output += [np.stack(mini_out, axis=0)]
return [np.concatenate(func_output, axis=1)] + self.step_states
self.merged_func = merged_func
self.merged_func.__name__ = utils.sanitize_name("_".join(
[type(op.process).__name__ for op in self.ops]))
def __init__(self, ops, signals, config):
super(GenericNeuronBuilder, self).__init__(ops, signals, config)
self.J_data = signals.combine([op.J for op in ops])
self.output_data = signals.combine([op.output for op in ops])
self.state_data = [signals.combine([op.states[i] for op in ops])
for i in range(len(ops[0].states))]
self.prev_result = []
def neuron_step_math(dt, J, *states): # pragma: no cover
output = None
J_offset = 0
state_offset = [0 for _ in states]
for op in ops:
# slice out the individual state vectors from the overall
# array
op_J = J[J_offset:J_offset + op.J.shape[0]]
J_offset += op.J.shape[0]
op_states = []
for j, s in enumerate(op.states):
op_states += [states[j][state_offset[j]:
state_offset[j] + s.shape[0]]]
state_offset[j] += s.shape[0]
# call step_math function
# note: `op_states` are views into `states`, which will
# be updated in-place
mini_out = []
for j in range(signals.minibatch_size):
# blank output variable
neuron_output = np.zeros(
op.output.shape, self.output_data.dtype)
op.neurons.step_math(dt, op_J[..., j], neuron_output,
*[s[..., j] for s in op_states])
mini_out += [neuron_output]
neuron_output = np.stack(mini_out, axis=-1)
# concatenate outputs
if output is None:
output = neuron_output
else:
output = np.concatenate((output, neuron_output),
axis=0)
return (output,) + states
self.neuron_step_math = neuron_step_math
self.neuron_step_math.__name__ = utils.sanitize_name(
"_".join([repr(op.neurons) for op in ops]))
def build_inputs(self, progress):
"""
Sets up the inputs in the model (which will be computed outside of
TensorFlow and fed in each simulation block).
Parameters
----------
progress : `.utils.ProgressBar`
Progress bar for input construction
"""
self.input_ph = {}
for n in progress(self.invariant_inputs):
if self.model.sig[n]["out"] in self.signals:
# set up a placeholder input for this node
self.input_ph[n] = tf.placeholder(
self.dtype, (None, n.size_out, self.minibatch_size),
name="%s_ph" % utils.sanitize_name(n))
def build_inputs(self, progress):
"""
Sets up the inputs in the model (which will be computed outside of
TensorFlow and fed in each simulation block).
Parameters
----------
progress : `.utils.ProgressBar`
Progress bar for input construction
"""
self.input_ph = {}
for n in progress(self.invariant_inputs):
if self.model.sig[n]["out"] in self.signals:
# set up a placeholder input for this node
self.input_ph[n] = tf.placeholder(
self.dtype, (None, n.size_out, self.minibatch_size),
name="%s_ph" % utils.sanitize_name(n))
def __init__(self, ops, signals, config):
super(GenericProcessBuilder, self).__init__(ops, signals, config)
self.input_data = (None if ops[0].input is None else
signals.combine([op.input for op in ops]))
self.output_data = signals.combine([op.output for op in ops])
self.output_shape = self.output_data.shape + (signals.minibatch_size,)
self.mode = "inc" if ops[0].mode == "inc" else "update"
self.prev_result = []
# build the step function for each process
self.step_fs = [[None for _ in range(signals.minibatch_size)]
for _ in ops]
# `merged_func` calls the step function for each process and
# combines the result
@utils.align_func(self.output_shape, self.output_data.dtype)
def merged_func(time, input): # pragma: no cover
if any(x is None for a in self.step_fs for x in a):
raise SimulationError(
"build_post has not been called for %s" % self)
input_offset = 0
func_output = []
for i, op in enumerate(ops):
if op.input is not None:
input_shape = op.input.shape[0]
func_input = input[input_offset:input_offset + input_shape]
input_offset += input_shape
mini_out = []
for j in range(signals.minibatch_size):
x = [] if op.input is None else [func_input[..., j]]
mini_out += [self.step_fs[i][j](*([time] + x))]
func_output += [np.stack(mini_out, axis=-1)]
return np.concatenate(func_output, axis=0)
self.merged_func = merged_func
self.merged_func.__name__ = utils.sanitize_name(
"_".join([type(op.process).__name__ for op in ops]))
def __init__(self, ops, signals, config):
super(GenericProcessBuilder, self).__init__(ops, signals, config)
self.input_data = (None if ops[0].input is None else signals.combine(
[op.input for op in ops]))
self.output_data = signals.combine([op.output for op in ops])
self.output_shape = self.output_data.shape + (signals.minibatch_size, )
self.mode = "inc" if ops[0].mode == "inc" else "update"
self.prev_result = []
# build the step function for each process
self.step_fs = [[None for _ in range(signals.minibatch_size)]
for _ in ops]
# `merged_func` calls the step function for each process and
# combines the result
@utils.align_func(self.output_shape, self.output_data.dtype)
def merged_func(time, input): # pragma: no cover
if any(x is None for a in self.step_fs for x in a):
raise SimulationError("build_post has not been called for %s" %
self)
input_offset = 0
func_output = []
for i, op in enumerate(ops):
if op.input is not None:
input_shape = op.input.shape[0]
func_input = input[input_offset:input_offset + input_shape]
input_offset += input_shape
mini_out = []
for j in range(signals.minibatch_size):
x = [] if op.input is None else [func_input[..., j]]
mini_out += [self.step_fs[i][j](*([time] + x))]
func_output += [np.stack(mini_out, axis=-1)]
return np.concatenate(func_output, axis=0)
self.merged_func = merged_func
self.merged_func.__name__ = utils.sanitize_name("_".join(
[type(op.process).__name__ for op in ops]))
def __init__(self, ops, signals, rng):
self.input_data = (None if ops[0].input is None else signals.combine(
[op.input for op in ops]))
self.output_data = signals.combine([op.output for op in ops])
self.output_shape = self.output_data.shape + (signals.minibatch_size, )
self.mode = "inc" if ops[0].mode == "inc" else "update"
self.prev_result = []
# build the step function for each process
step_fs = [[
op.process.make_step(
op.input.shape if op.input is not None else (0, ),
op.output.shape, signals.dt_val, op.process.get_rng(rng))
for _ in range(signals.minibatch_size)
] for op in ops]
# `merged_func` calls the step function for each process and
# combines the result
@utils.align_func(self.output_shape, self.output_data.dtype)
def merged_func(time, input): # pragma: no cover
input_offset = 0
func_output = []
for i, op in enumerate(ops):
if op.input is not None:
input_shape = op.input.shape[0]
func_input = input[input_offset:input_offset + input_shape]
input_offset += input_shape
mini_out = []
for j in range(signals.minibatch_size):
x = [] if op.input is None else [func_input[..., j]]
mini_out += [step_fs[i][j](*([time] + x))]
func_output += [np.stack(mini_out, axis=-1)]
return np.concatenate(func_output, axis=0)
self.merged_func = merged_func
self.merged_func.__name__ = utils.sanitize_name("_".join(
[type(op.process).__name__ for op in ops]))
def __init__(self, log_dir, sim, objects):
super().__init__()
self.sim = sim
# we do all the summary writing in eager mode, so that it will be executed
# as the callback is called
with context.eager_mode():
self.writer = tf.summary.create_file_writer(log_dir)
self.summaries = []
for obj in objects:
if isinstance(
obj,
(nengo.Ensemble, nengo.ensemble.Neurons, nengo.Connection)):
if isinstance(obj, nengo.Ensemble):
param = "encoders"
name = "Ensemble_%s" % obj.label
elif isinstance(obj, nengo.ensemble.Neurons):
param = "bias"
name = "Ensemble.neurons_%s" % obj.ensemble.label
elif isinstance(obj, nengo.Connection):
if not compat.conn_has_weights(obj):
raise ValidationError(
"Connection '%s' does not have any weights to log"
% obj,
"objects",
)
param = "weights"
name = "Connection_%s" % obj.label
self.summaries.append(
(utils.sanitize_name("%s_%s" % (name, param)), obj, param))
else:
raise ValidationError(
"Unknown summary object %s; should be an Ensemble, Neurons, or "
"Connection" % obj,
"objects",
)
def name_scope(self, ops):
"""Returns a new TensorFlow name scope for the given ops."""
return self.graph.name_scope(
utils.sanitize_name(Builder.builders[type(ops[0])].__name__))
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.summary.scalar("loss",
tf.reduce_sum([
tf.reduce_sum(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.summary.scalar(utils.sanitize_name(
"Probe_%s_loss" % p.label),
tf.reduce_sum(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.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.summary.merge(summary_ops), (None if len(inits) == 0 else
inits)
def test_sanitize_name():
assert utils.sanitize_name(0) == "0"
assert utils.sanitize_name("a b") == "a_b"
assert utils.sanitize_name("a:b") == "a_b"
assert utils.sanitize_name(r"Aa0.-/\,?^&*") == r"Aa0.-/"
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 positional arguments for function
spec = inspect.getfullargspec(out)
nargs = len(spec.args)
if spec.defaults is not None:
# don't count keyword arguments
nargs -= len(spec.defaults)
if inspect.ismethod(out) or not inspect.isroutine(out):
# don't count self argument for methods or callable classes
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.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.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.GraphKeys.GLOBAL_VARIABLES,
tf.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.variables_initializer(new_vars)
if len(new_vars) > 0 else None)
self.outputs[key] = output_vals
return output_vals, new_vars_init
def build_pre(self, signals, config):
super().build_pre(signals, config)
self.J_data = signals.combine([op.J for op in self.ops])
self.output_data = signals.combine([op.output for op in self.ops])
state_keys = compat.neuron_state(self.ops[0]).keys()
self.state_data = [
signals.combine([compat.neuron_state(op)[key] for op in self.ops])
for key in state_keys
]
self.prev_result = []
def neuron_step(dt, J, *states): # pragma: no cover (runs in TF)
output = None
J_offset = 0
state_offset = [0 for _ in states]
for op in self.ops:
# slice out the individual state vectors from the overall
# array
op_J = J[:, J_offset:J_offset + op.J.shape[0]]
J_offset += op.J.shape[0]
op_states = []
for j, key in enumerate(state_keys):
s = compat.neuron_state(op)[key]
op_states += [
states[j][:,
state_offset[j]:state_offset[j] + s.shape[0]]
]
state_offset[j] += s.shape[0]
# call neuron step function
# note: `op_states` are views into `states`, which will
# be updated in-place
mini_out = []
for j in range(signals.minibatch_size):
# blank output variable
neuron_output = np.zeros(op.output.shape,
self.output_data.dtype)
compat.neuron_step(
op,
dt,
op_J[j],
neuron_output,
dict(zip(state_keys, [s[j] for s in op_states])),
)
mini_out.append(neuron_output)
neuron_output = np.stack(mini_out, axis=0)
# concatenate outputs
if output is None:
output = neuron_output
else:
output = np.concatenate((output, neuron_output), axis=1)
return (output, ) + states
self.neuron_step = neuron_step
self.neuron_step.__name__ = utils.sanitize_name("_".join(
[repr(op.neurons) for op in self.ops]))
def build(self, progress):
"""
Constructs a new graph to simulate the model.
progress : :class:`.utils.ProgressBar`
Progress bar for construction stage
"""
self.signals = signals.SignalDict(self.sig_map, self.dtype,
self.minibatch_size)
self.target_phs = {}
self.losses = {}
self.optimizers = {}
# make sure indices are loaded for all probe signals (they won't
# have been loaded if this signal is only accessed as part of a
# larger block during the simulation)
for p in self.model.probes:
probe_sig = self.model.sig[p]["in"]
if probe_sig in self.sig_map:
self.sig_map[probe_sig].load_indices()
# create this constant once here so we don't end up creating a new
# dt constant in each operator
self.signals.dt = tf.constant(self.dt, self.dtype)
self.signals.dt_val = self.dt # store the actual value as well
# variable to track training step
with tf.device("/cpu:0"):
with tf.variable_scope("misc_vars", reuse=False):
self.training_step = tf.get_variable(
"training_step",
initializer=tf.constant_initializer(0),
dtype=tf.int64,
shape=(),
trainable=False)
self.training_step_inc = tf.assign_add(self.training_step, 1)
# 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.placeholder(v.dtype, v.shape)
if trainable:
with tf.variable_scope("trainable_vars", reuse=False):
var = tf.get_variable(name, initializer=ph, trainable=True)
else:
with tf.variable_scope("local_vars", reuse=False):
var = tf.get_local_variable(name,
initializer=ph,
trainable=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
sub = progress.sub("pre-build stage")
self.op_builds = {}
for ops in sub(self.plan):
with self.graph.name_scope(
utils.sanitize_name(builder.Builder.builders[type(
ops[0])].__name__)):
builder.Builder.pre_build(ops, self.signals, self.op_builds)
# build stage
sub = progress.sub("unrolled step ops")
self.build_loop(sub)
# ops for initializing variables (will be called by simulator)
trainable_vars = tf.trainable_variables() + [self.training_step]
self.trainable_init_op = tf.variables_initializer(trainable_vars)
self.local_init_op = tf.local_variables_initializer()
self.global_init_op = tf.variables_initializer(
[v for v in tf.global_variables() if v not in trainable_vars])
self.constant_init_op = tf.variables_initializer(
tf.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)
def build_step(self):
"""
Build the operators that execute a single simulation timestep
into the graph.
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., :class:`~nengo: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
"""
# build operators
side_effects = []
# 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
for ops in self.plan:
with self.graph.name_scope(
utils.sanitize_name(builder.Builder.builders[type(
ops[0])].__name__)):
outputs = builder.Builder.build(ops, self.signals,
self.op_builds)
if outputs is not None:
side_effects += outputs
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.sig_map:
# 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.sig_map[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
def name_scope(self, ops):
"""Returns a new TensorFlow name scope for the given ops."""
return tf.name_scope(
utils.sanitize_name(Builder.builders[type(ops[0])].__name__)
)
def __init__(self, model, dt, unroll_simulation, minibatch_size, device,
progress, seed):
super().__init__(
name="TensorGraph",
dynamic=False,
trainable=not config.get_setting(model, "inference_only", False),
dtype=config.get_setting(model, "dtype", "float32"),
batch_size=minibatch_size,
)
self.model = model
self.dt = dt
self.unroll = unroll_simulation
self.use_loop = config.get_setting(model, "use_loop", True)
self.minibatch_size = minibatch_size
self.device = device
self.seed = seed
self.inference_only = not self.trainable
self.signals = signals.SignalDict(self.dtype, self.minibatch_size)
# 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))
# 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, 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.default_simplifications,
)
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)
# generate unique names for layer inputs/outputs
# this follows the TensorFlow unique naming scheme, so if multiple objects are
# created with the same name, they will be named like name, NAME_1, name_2
# (note: case insensitive)
self.io_names = {}
name_count = defaultdict(int)
for obj in list(self.invariant_inputs.keys()) + self.model.probes:
name = (type(obj).__name__.lower()
if obj.label is None else utils.sanitize_name(obj.label))
key = name.lower()
if name_count[key] > 0:
name += "_%d" % name_count[key]
self.io_names[obj] = name
name_count[key] += 1
logger.info("Optimized plan length: %d", len(self.plan))
logger.info(
"Number of base arrays: (%s, %d), (%s, %d), (%s, %d)",
*tuple((k, len(x)) for k, x in self.base_arrays_init.items()),
)
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 positional arguments for function
spec = inspect.getfullargspec(out)
nargs = len(spec.args)
if spec.defaults is not None:
# don't count keyword arguments
nargs -= len(spec.defaults)
if inspect.ismethod(out) or not inspect.isroutine(out):
# don't count self argument for methods or callable classes
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.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.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.GraphKeys.GLOBAL_VARIABLES,
tf.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.variables_initializer(new_vars)
if len(new_vars) > 0 else None)
self.outputs[key] = output_vals
return output_vals, new_vars_init
def test_sanitize_name():
assert utils.sanitize_name(0) == "0"
assert utils.sanitize_name("a b") == "a_b"
assert utils.sanitize_name("a:b") == "a_b"
assert utils.sanitize_name(r"Aa0.-/\,?^&*") == r"Aa0.-/"
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.summary.scalar(
"loss", tf.reduce_sum([tf.reduce_sum(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.summary.scalar(
utils.sanitize_name("Probe_%s_loss" % p.label),
tf.reduce_sum(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.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.summary.merge(summary_ops), (None if len(inits) == 0 else
inits)
