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)