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)
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(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)