def format_transform(size, transform): if is_transform_type(transform, "NoTransform"): transform = np.array(1.0) elif is_transform_type(transform, "Dense"): transform = transform.init else: raise NotImplementedError( "Mergeable transforms must be Dense; " "set remove_passthrough=False" ) if not isinstance(transform, np.ndarray): raise NotImplementedError( "Mergeable transforms must be specified as Numpy arrays, " "not distributions. Set `remove_passthrough=False`." ) if transform.ndim == 0: # scalar transform = np.eye(size) * transform elif transform.ndim != 2: raise BuildError( f"{node}: Unhandled transform shape: {transform.shape}" ) return transform
def build_host_to_learning_rule(model, conn): if not is_transform_type(conn.transform, ("Dense", "NoTransform")): # TODO: What needs to be done to support this? It looks like it should just work raise BuildError( f"{conn}: nengo-loihi does not yet support " f"'{type(conn.transform).__name__}' transforms on host to chip " "learning rule connections" ) dim = conn.size_out host = model.host_model(base_obj(conn.pre)) logger.debug("Creating HostSendNode for %s", conn) send = HostSendNode( dim, label=None if conn.label is None else "%s_send" % conn.label, add_to_container=False, ) host.build(send) pre2send = Connection( conn.pre, send, function=conn.function, solver=conn.solver, eval_points=conn.eval_points, scale_eval_points=conn.scale_eval_points, synapse=conn.synapse, transform=conn.transform, label=conn.label, add_to_container=False, ) model.host2chip_pes_senders[send] = model.needs_sender[conn.post_obj] _inherit_seed(host, pre2send, model, conn) host.build(pre2send)
def build_chip_to_host(model, conn): if not is_transform_type(conn.transform, ("Dense", "NoTransform")): raise BuildError( f"{conn}: nengo-loihi does not yet support " f"'{type(conn.transform).__name__}' transforms on chip to host connections" ) rng = np.random.RandomState(model.seeds[conn]) dim = conn.size_out host = model.host_model(base_obj(conn.post)) logger.debug("Creating HostReceiveNode for %s", conn) receive = HostReceiveNode( dim, label=None if conn.label is None else "%s_receive" % conn.label, add_to_container=False, ) host.build(receive) receive2post = Connection( receive, conn.post, synapse=conn.synapse, label=None if conn.label is None else "%s_host" % conn.label, add_to_container=False, ) _inherit_seed(host, receive2post, model, conn) host.build(receive2post) logger.debug("Creating Probe for %s", conn) transform = sample_transform(conn, rng=rng) probe = NengoProbe( conn.pre, synapse=None, solver=conn.solver, add_to_container=False ) model.chip2host_params[probe] = dict( learning_rule_type=conn.learning_rule_type, function=conn.function, eval_points=conn.eval_points, scale_eval_points=conn.scale_eval_points, transform=transform, label=None if conn.label is None else "%s_probe" % conn.label, ) model.chip2host_receivers[probe] = receive _inherit_seed(model, probe, model, conn) model.builder.build(model, probe) if conn.learning_rule_type is not None: if not isinstance(conn.pre_obj, Ensemble): raise NotImplementedError( "Learning rule presynaptic object must be an Ensemble " "(got %r)" % type(conn.pre_obj).__name__ ) model.needs_sender[conn.learning_rule] = PESModulatoryTarget(probe)
def build_decoders(model, conn, rng, sampled_transform): # Copied from Nengo, except where noted below encoders = model.params[conn.pre_obj].encoders gain = model.params[conn.pre_obj].gain bias = model.params[conn.pre_obj].bias eval_points = get_eval_points(model, conn, rng) targets = get_targets(conn, eval_points) if conn.solver.weights and not conn.solver.compositional: # solver is solving for the whole weight matrix, so apply # transform/encoders to targets # CHANGE: backwards compatibility with nengo<=2.8.0 # if not isinstance(conn.transform, Dense): # raise BuildError( # "Non-compositional solvers only work with Dense transforms") # transform = conn.transform.sample(rng=rng) # targets = np.dot(targets, transform.T) if not is_transform_type(conn.transform, "Dense"): # pragma: no cover raise BuildError( f"{conn}: Non-compositional solvers only work with Dense transforms" ) targets = np.dot(targets, sampled_transform.T) # weight solvers only allowed on ensemble->ensemble connections assert isinstance(conn.post_obj, Ensemble) post_enc = model.params[conn.post_obj].scaled_encoders targets = np.dot(targets, post_enc.T[conn.post_slice]) x = np.dot(eval_points, encoders.T / conn.pre_obj.radius) # CHANGE: we pass `dt` to `solve_for_decoders`, # and do not support the decoder cache. # wrapped_solver = (model.decoder_cache.wrap_solver(solve_for_decoders) # if model.seeded[conn] else solve_for_decoders) # decoders, solver_info = wrapped_solver( # conn, gain, bias, x, targets, rng=rng) decoders, solver_info = solve_for_decoders( conn, gain, bias, x, targets, rng=rng, dt=model.dt ) return eval_points, decoders.T, solver_info
def build_conv2d_connection(model, transform, conn): assert is_transform_type(transform, ("Convolution", "ConvolutionTranspose")) if transform.dimensions != 2: raise NotImplementedError("nengo-loihi only supports 2D convolution") # Create random number generator rng = np.random.RandomState(model.seeds[conn]) pre_obj = model.objs[conn.pre_obj]["out"] post_obj = model.objs[conn.post_obj]["in"] assert isinstance(pre_obj, (LoihiInput, LoihiBlock)) assert isinstance(post_obj, LoihiBlock) tau_s = 0.0 if isinstance(conn.synapse, nengo.synapses.Lowpass): tau_s = conn.synapse.tau elif conn.synapse is not None: raise NotImplementedError("Cannot handle non-Lowpass synapses") # --- pre assert isinstance(conn.pre_obj, (Neurons, ChipReceiveNeurons)) kernel = transform.sample(rng=rng) input_shape = transform.input_shape # Account for nengo spike height of 1/dt kernel = kernel / model.dt if isinstance(conn.pre_obj, ChipReceiveNeurons): neuron_type = conn.pre_obj.neuron_type elif isinstance(conn.pre_obj, Neurons): neuron_type = conn.pre_obj.ensemble.neuron_type if neuron_type is not None and hasattr(neuron_type, "amplitude"): kernel = kernel * neuron_type.amplitude # --- post assert isinstance(conn.post_obj, Neurons) assert conn.post_slice == slice(None) gain = model.params[conn.post_obj.ensemble].gain if not np.all(gain == gain[0]): # Cannot fold gains into kernel, result would not be convolutional. # Therefore, Loihi does not support this if we want to share weights. raise ValidationError( "All neurons targeted by a Convolution connection must " "have the same gain", "gain", obj=conn.post_obj.ensemble, ) kernel = kernel * gain[0] kernel = kernel.astype(nengo.rc.float_dtype) pop_type = model.config[conn].pop_type new_transform = copy.copy(transform) type(new_transform).init.data[new_transform] = kernel weights, indices, axon_to_weight_map, offsets = conv2d_loihi_weights(new_transform) synapse = Synapse(np.prod(input_shape.spatial_shape), label="conv2d_weights") synapse.set_population_weights( weights, indices, axon_to_weight_map, offsets, pop_type=pop_type ) post_obj.add_synapse(synapse) model.objs[conn]["weights"] = synapse if synapse.atom_bits_extra() > 0: warnings.warn( "Using more than 32 'populations' (e.g. convolutional filters) with " "`pop_type=16` axons has not yet been implemented in NxSDK. This feature " "is therefore emulator-only." ) target_axons = -np.ones(pre_obj.n_neurons, dtype=np.int32) target_axons[conn.pre_slice] = pixel_idxs(input_shape) atoms = np.zeros(pre_obj.n_neurons, dtype=np.int32) atoms[conn.pre_slice] = channel_idxs(input_shape) ax = Axon(np.prod(input_shape.spatial_shape), label="conv2d_weights") ax.target = synapse ax.set_compartment_axon_map(target_axons, atoms=atoms) pre_obj.add_axon(ax) post_obj.compartment.configure_filter(tau_s, dt=model.dt) model.params[conn] = BuiltConnection( eval_points=None, solver_info=None, transform=None, weights=kernel )
def build_chip_connection(model, conn): if is_transform_type(conn.transform, ("NoTransform", "Dense", "Sparse")): return build_full_chip_connection(model, conn) else: model.build(conn.transform, conn)
def build_host_to_chip(model, conn): rng = np.random.RandomState(model.seeds[conn]) host = model.host_model(base_obj(conn.pre)) if is_transform_type(conn.transform, ("Convolution", "ConvolutionTranspose")): raise BuildError( f"{conn}: Conv2D transforms not supported for off-chip to " "on-chip connections where `pre` is not a Neurons object." ) elif not is_transform_type(conn.transform, ("Dense", "NoTransform")): raise BuildError( f"{conn}: nengo-loihi does not yet support " f"'{type(conn.transform).__name__}' transforms on host to chip connections" ) # Scale the input spikes based on the radius of the target ensemble weights = sample_transform(conn, rng=rng) if isinstance(conn.post_obj, Ensemble): weights = weights / conn.post_obj.radius if is_transform_type(conn.transform, "NoTransform"): transform = weights # weights are 1 / (post ensemble radius), if applicable else: # copy the Transform information, setting `init` to the sampled weights transform = copy.copy(conn.transform) type(transform).init.data[transform] = weights if isinstance(conn.post_obj, Neurons): # we don't have encoders, and the transform could have large output, # so do it on the chip host_transform = 1.0 chip_transform = transform dim = conn.size_mid else: # we have encoders on the chip, so do the transform off-chip host_transform = transform chip_transform = 1.0 dim = conn.size_out logger.debug("Creating ChipReceiveNode for %s", conn) receive = ChipReceiveNode( dim * 2, size_out=dim, label=None if conn.label is None else "%s_node" % conn.label, add_to_container=False, ) model.builder.build(model, receive) receive2post = Connection( receive, conn.post, transform=chip_transform, synapse=model.decode_tau, label=None if conn.label is None else "%s_chip" % conn.label, add_to_container=False, ) _inherit_seed(model, receive2post, model, conn) _inherit_config(model, receive2post, model, conn) build_chip_connection(model, receive2post) logger.debug("Creating DecodeNeuron ensemble for %s", conn) ens = model.node_neurons.get_ensemble(dim, add_to_container=False) ens.label = None if conn.label is None else "%s_ens" % conn.label _inherit_seed(host, ens, model, conn) host.build(ens) model.connection_decode_neurons[conn] = ens pre2ens = Connection( conn.pre, ens, function=conn.function, solver=conn.solver, eval_points=conn.eval_points, scale_eval_points=conn.scale_eval_points, synapse=conn.synapse, transform=host_transform, label=None if conn.label is None else "%s_enc" % conn.label, add_to_container=False, ) _inherit_seed(host, pre2ens, model, conn) host.build(pre2ens) logger.debug("Creating HostSendNode for %s", conn) send = HostSendNode( dim * 2, label=None if conn.label is None else "%s_send" % conn.label, add_to_container=False, ) host.build(send) ensneurons2send = Connection( ens.neurons, send, synapse=None, label=None if conn.label is None else "%s_host" % conn.label, add_to_container=False, ) _inherit_seed(host, ensneurons2send, model, conn) model.host2chip_senders[send] = receive host.build(ensneurons2send)
def conv2d_loihi_weights(transform): # noqa: C901 assert ( transform.channels_last == transform.input_shape.channels_last ), "Transforms that switch the channel position not yet implemented" transpose = is_transform_type(transform, "ConvolutionTranspose") input_rows, input_cols = transform.input_shape.spatial_shape n_channels = transform.input_shape.n_channels output_rows, output_cols = transform.output_shape.spatial_shape n_filters = transform.n_filters n_compartments = output_rows * output_cols * n_filters kernel_rows, kernel_cols = transform.kernel_size row_stride, col_stride = transform.strides kernel = transform.init assert isinstance(kernel, np.ndarray), "Should already have been sampled" assert kernel.shape == (kernel_rows, kernel_cols, n_channels, n_filters) # tranpose kernel to (in_channels, rows, cols, out_channels) kernel = np.transpose(kernel, (2, 0, 1, 3)) if not transpose: # flip weights to do correlation kernel = kernel[:, ::-1, ::-1, :] # compute number of used input pixels if not transpose: ri_max = (output_rows - 1) * row_stride + 1 rj_max = (output_cols - 1) * col_stride + 1 else: # compute number of used output pixels ri_max, rj_max = transform.output_shape.spatial_shape # --- determine padding pad_i, pad_j = 0, 0 if transform.padding == "same": if transpose: # these paddings are based off the method used in # `nengo._vendor.npconv2d`, to ensure we perform the same output_rows_min = (input_rows - 1) * row_stride + 1 output_cols_min = (input_cols - 1) * col_stride + 1 pad_i = min( max(output_rows + kernel_rows - 1 - output_rows_min, 0), (kernel_rows - 1) * 2, ) pad_j = min( max(output_cols + kernel_cols - 1 - output_cols_min, 0), (kernel_cols - 1) * 2, ) # use floor instead of the `ceil` used by `npconv2d.conv2d_gradx`, since # this padding is applied to the output where the kernel is flipped pad_i, pad_j = pad_i // 2, pad_j // 2 pad_i, pad_j = -pad_i, -pad_j else: # these paddings are based off the method used in # `nengo._vendor.npconv2d`, to ensure we perform the same pad_i = max( (output_rows - 1) * row_stride + kernel_rows - input_rows, 0) pad_j = max( (output_cols - 1) * col_stride + kernel_cols - input_cols, 0) pad_i, pad_j = pad_i // 2, pad_j // 2 # --- determine weights and indices weights = [] indices = [] # compartment offset (aka. compartment base) for each axon offsets = np.zeros(input_rows * input_cols, dtype=np.int32) axon_to_weight_map = np.zeros(input_rows * input_cols, dtype=np.int32) weights_map = {} for i, j in itertools.product(range(input_rows), range(input_cols)): ij = i * input_cols + j if transpose: # compartment indices that this input axon would map to if mode == 'valid' ri0 = i * row_stride + pad_i rj0 = j * col_stride + pad_j else: # unstrided compartment indices that this input axon would map to # if strides == 1 and mode == 'full' ri0 = i + pad_i + 1 - kernel_rows rj0 = j + pad_j + 1 - kernel_cols ri = np.arange(ri0, ri0 + kernel_rows) rj = np.arange(rj0, rj0 + kernel_cols) wmask_i = (ri >= 0) & (ri < ri_max) wmask_j = (rj >= 0) & (rj < rj_max) if transpose: assert wmask_i.sum() > 0 and wmask_j.sum() > 0 else: wmask_i &= ri % row_stride == 0 wmask_j &= rj % col_stride == 0 if wmask_i.sum() == 0 or wmask_j.sum() == 0: # this axon is not needed, so indicate this in offsets and skip offsets[ij] = -1 continue yi0, yj0 = ri[wmask_i][0], rj[wmask_j][0] if not transpose: yi0 = yi0 // row_stride yj0 = yj0 // col_stride yij0 = yi0 * output_cols + yj0 offset = yij0 * n_filters if transform.channels_last else yij0 # There is currently an upper limit on the axon compartment offset of 256. # To work around this, we split the offset into two parts, and make extra sets # of redundant weights with part of the offset in the indices, as needed. axon_offset = offset % 256 index_offset = offset - axon_offset offsets[ij] = axon_offset weight_key = (tuple(wmask_i), tuple(wmask_j), index_offset) if weight_key not in weights_map: w = kernel[:, wmask_i[:, None] * wmask_j, :] assert w.shape == (n_channels, wmask_i.sum() * wmask_j.sum(), n_filters) # --- determine indices # channel inds are zero, since we use same indices for each channel channel_inds = np.zeros(n_channels, dtype=np.int32) row_inds = np.arange(wmask_i.sum(), dtype=np.int32) col_inds = np.arange(wmask_j.sum(), dtype=np.int32) filter_inds = np.arange(n_filters, dtype=np.int32) order = [channel_inds, row_inds, col_inds, filter_inds] shape = [n_channels, output_rows, output_cols, n_filters] if not transform.channels_last: # move filters (aka. output channels) before rows/cols w = np.transpose(w, (0, 2, 1)) order = [order[i] for i in (0, 3, 1, 2)] shape = [shape[i] for i in (0, 3, 1, 2)] n = len(shape) strides = [ np.prod(shape[i + 1:], dtype=np.int32) for i in range(n) ] # inds[i_0,...,i_{n-1}] = sum_{k=0}^{n-1} strides[k] * order[k][i_k] strided_inds = [ stride * ind.reshape([-1] + [1] * (n - 1 - k)) for k, (ind, stride) in enumerate(zip(order, strides)) ] inds = sum([index_offset] + strided_inds) weights_map[weight_key] = len(weights) weights.append(w.reshape(n_channels, -1)) indices.append(inds.reshape(n_channels, -1)) axon_to_weight_map[ij] = weights_map[weight_key] # check that offset (compartment base) plus index points to a valid compartment inds = indices[axon_to_weight_map[ij]] assert (offsets[ij] + inds < n_compartments).all() return weights, indices, axon_to_weight_map, offsets