import logging
import warnings
from nengo.builder.processes import SimProcess
from nengo.exceptions import SimulationError
from nengo.synapses import Lowpass, LinearFilter
from nengo.utils.filter_design import (cont2discrete, tf2ss, ss2tf,
BadCoefficients)
import numpy as np
import tensorflow as tf
from tensorflow.python.ops import gen_sparse_ops
from nengo_dl import utils
from nengo_dl.builder import Builder, OpBuilder
logger = logging.getLogger(__name__)
[docs]@Builder.register(SimProcess)
class SimProcessBuilder(OpBuilder):
"""Builds a group of :class:`~nengo:nengo.builder.processes.SimProcess`
operators.
Calls the appropriate sub-build class for the different process types.
Attributes
----------
TF_PROCESS_IMPL : list of :class:`~nengo:nengo.Process`
The process types that have a custom implementation
"""
TF_PROCESS_IMPL = (Lowpass, LinearFilter)
def __init__(self, ops, signals):
super(SimProcessBuilder, self).__init__(ops, signals)
logger.debug("process %s", [op.process for op in ops])
logger.debug("input %s", [op.input for op in ops])
logger.debug("output %s", [op.output for op in ops])
logger.debug("t %s", [op.t for op in ops])
# if we have a custom tensorflow implementation for this process type,
# then we build that. otherwise we'll execute the process step
# function externally (using `tf.py_func`), so we just need to set up
# the inputs/outputs for that.
if isinstance(ops[0].process, self.TF_PROCESS_IMPL):
# note: we do this two-step check (even though it's redundant) to
# make sure that TF_PROCESS_IMPL is kept up to date
if type(ops[0].process) == Lowpass:
self.built_process = LowpassBuilder(ops, signals)
elif isinstance(ops[0].process, LinearFilter):
self.built_process = LinearFilterBuilder(ops, signals)
else:
self.built_process = GenericProcessBuilder(ops, signals)
[docs] def build_step(self, signals):
self.built_process.build_step(signals)
[docs] def build_post(self, ops, signals, sess, rng):
if isinstance(self.built_process, GenericProcessBuilder):
self.built_process.build_post(ops, signals, sess, rng)
[docs]class GenericProcessBuilder(OpBuilder):
"""Builds all process types for which there is no custom TensorFlow
implementation.
Notes
-----
These will be executed as native Python functions, requiring execution to
move in and out of TensorFlow. This can significantly slow down the
simulation, so any performance-critical processes should consider
adding a custom TensorFlow implementation for their type instead.
"""
def __init__(self, ops, signals):
super(GenericProcessBuilder, self).__init__(ops, signals)
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]))
[docs] def build_step(self, signals):
input = ([] if self.input_data is None
else signals.gather(self.input_data))
# note: we need to make sure that the previous call to this function
# has completed before the next starts, since we don't know that the
# functions are thread safe
with tf.control_dependencies(self.prev_result), tf.device("/cpu:0"):
result = tf.py_func(
self.merged_func, [signals.time, input],
self.output_data.dtype, name=self.merged_func.__name__)
result.set_shape(self.output_shape)
self.prev_result = [result]
signals.scatter(self.output_data, result, mode=self.mode)
[docs] def build_post(self, ops, signals, sess, rng):
for i, op in enumerate(ops):
for j in range(signals.minibatch_size):
self.step_fs[i][j] = 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))
[docs]class LowpassBuilder(OpBuilder):
"""Build a group of :class:`~nengo:nengo.Lowpass` synapse operators."""
def __init__(self, ops, signals):
super(LowpassBuilder, self).__init__(ops, signals)
self.input_data = signals.combine([op.input for op in ops])
self.output_data = signals.combine([op.output for op in ops])
nums = []
dens = []
for op in ops:
if op.process.tau <= 0.03 * signals.dt_val:
num = 1
den = 0
else:
num, den, _ = cont2discrete((op.process.num, op.process.den),
signals.dt_val, method="zoh")
num = num.flatten()
num = num[1:] if num[0] == 0 else num
assert len(num) == 1
num = num[0]
assert len(den) == 2
den = den[1]
nums += [num] * op.input.shape[0]
dens += [den] * op.input.shape[0]
nums = np.asarray(nums)
while nums.ndim < len(self.input_data.full_shape):
nums = np.expand_dims(nums, -1)
# note: applying the negative here
dens = -np.asarray(dens)
while dens.ndim < len(self.input_data.full_shape):
dens = np.expand_dims(dens, -1)
# need to manually broadcast for scatter_mul
# dens = np.tile(dens, (1, signals.minibatch_size))
self.nums = tf.constant(nums, dtype=self.output_data.dtype)
self.dens = tf.constant(dens, dtype=self.output_data.dtype)
# create a variable to represent the internal state of the filter
# self.state_sig = signals.make_internal(
# "state", self.output_data.shape)
[docs] def build_step(self, signals):
# signals.scatter(self.output_data, self.dens, mode="mul")
# input = signals.gather(self.input_data)
# signals.scatter(self.output_data, self.nums * input, mode="inc")
input = signals.gather(self.input_data)
output = signals.gather(self.output_data)
signals.scatter(self.output_data,
self.dens * output + self.nums * input)
# method using _step
# note: this build_step function doesn't use _step for efficiency
# reasons (we can avoid an extra scatter by reusing the output signal
# as the state signal)
# input = signals.gather(self.input_data)
# prev_state = signals.gather(self.state_sig)
# new_state = self.dens * prev_state + self.nums * input
# signals.scatter(self.state_sig, new_state)
# signals.scatter(self.output_data, new_state)
[docs]class LinearFilterBuilder(OpBuilder):
"""Build a group of :class:`~nengo:nengo.LinearFilter` synapse
operators."""
def __init__(self, ops, signals):
super(LinearFilterBuilder, self).__init__(ops, signals)
self.input_data = signals.combine([op.input for op in ops])
self.output_data = signals.combine([op.output for op in ops])
self.n_ops = len(ops)
self.signal_d = ops[0].input.shape[0]
As = []
Cs = []
Ds = []
# compute the A/C/D matrices for each operator
for op in ops:
A, B, C, D = tf2ss(op.process.num, op.process.den)
if op.process.analog:
# convert to discrete system
A, B, C, D, _ = cont2discrete((A, B, C, D), signals.dt_val,
method="zoh")
# convert to controllable form
num, den = ss2tf(A, B, C, D)
if op.process.analog:
# add shift
num = np.concatenate((num, [[0]]), axis=1)
with warnings.catch_warnings():
# ignore the warning about B, since we aren't using it anyway
warnings.simplefilter("ignore", BadCoefficients)
A, _, C, D = tf2ss(num, den)
As.append(A)
Cs.append(C[0])
Ds.append(D.item())
self.state_d = sum(x.shape[0] for x in Cs)
# build a sparse matrix containing the A matrices as blocks
# along the diagonal
sparse_indices = []
corner = np.zeros(2, dtype=np.int64)
for A in As:
idxs = np.reshape(np.dstack(np.meshgrid(
np.arange(A.shape[0]), np.arange(A.shape[1]),
indexing="ij")), (-1, 2))
idxs += corner
corner += A.shape
sparse_indices += [idxs]
sparse_indices = np.concatenate(sparse_indices, axis=0)
self.A = tf.constant(np.concatenate(As, axis=0).flatten(),
dtype=signals.dtype)
self.A_indices = tf.constant(sparse_indices, dtype=(
tf.int32 if np.all(sparse_indices < np.iinfo(np.int32).max)
else tf.int64))
self.A_shape = tf.constant(corner, dtype=tf.int64)
if np.allclose(Cs, 0):
self.C = None
else:
# add empty dimension for broadcasting
self.C = tf.constant(np.concatenate(Cs)[:, None],
dtype=signals.dtype)
if np.allclose(Ds, 0):
self.D = None
else:
# add empty dimension for broadcasting
self.D = tf.constant(np.asarray(Ds)[:, None], dtype=signals.dtype)
self.offsets = tf.expand_dims(
tf.range(0, len(ops) * As[0].shape[0], As[0].shape[0]),
axis=1)
# create a variable to represent the internal state of the filter
self.state_sig = signals.make_internal(
"state", (self.state_d, signals.minibatch_size * self.signal_d),
minibatched=False)
[docs] def build_step(self, signals):
input = signals.gather(self.input_data)
input = tf.reshape(input, (self.n_ops, -1))
state = signals.gather(self.state_sig)
# compute output
if self.C is None:
output = tf.zeros_like(input)
else:
output = state * self.C
output = tf.reshape(
output,
(self.n_ops, -1, signals.minibatch_size * self.signal_d))
output = tf.reduce_sum(output, axis=1)
if self.D is not None:
output += self.D * input
signals.scatter(self.output_data, output)
# update state
if tf.__version__ < "1.7.0":
mat_mul = gen_sparse_ops._sparse_tensor_dense_mat_mul
else:
mat_mul = gen_sparse_ops.sparse_tensor_dense_mat_mul
r = mat_mul(self.A_indices, self.A, self.A_shape, state)
with tf.control_dependencies([output]):
state = r + tf.scatter_nd(self.offsets, input,
self.state_sig.shape)
# TODO: tensorflow does not yet support sparse_tensor_dense_add
# on the GPU
# state = gen_sparse_ops._sparse_tensor_dense_add(
# self.offsets, input, self.state_sig.shape, r)
state.set_shape(self.state_sig.shape)
signals.mark_gather(self.input_data)
signals.mark_gather(self.state_sig)
signals.scatter(self.state_sig, state)