Source code for nengo_extras.cuda_convnet

from __future__ import absolute_import

import collections
import os

import nengo
import numpy as np

from .compat import cmp, pickle_load
from .deepnetworks import SequentialNetwork


def load_model_pickle(loadfile):
    loadfile = os.path.expanduser(loadfile)
    with open(loadfile, 'rb') as f:
        return pickle_load(f)


def layer_depths(layers):
    depths = {}

    def get_depth(name):
        if name in depths:
            return depths[name]

        inputs = layers[name].get('inputs', [])
        depth = max(get_depth(i) for i in inputs) + 1 if len(inputs) > 0 else 0
        depths[name] = depth
        return depth

    for name in layers:
        get_depth(name)

    return depths


def sort_layers(layers, depths=None, cycle_check=True):
    depths = layer_depths(layers) if depths is None else depths

    def compare(a, b):
        da, db = depths[a], depths[b]
        return cmp(a, b) if da == db else cmp(da, db)

    snames = sorted(layers, cmp=compare)

    if cycle_check:
        def compare_cycle(a, b):
            ainb = a in layers[b].get('inputs', [])
            bina = b in layers[a].get('inputs', [])
            assert not (ainb and bina), "Cycle in graph"
            return -1 if ainb else 1 if bina else 0

        if any(compare_cycle(snames[i], snames[j]) < 0
               for i in range(1, len(snames)) for j in range(i)):
            raise ValueError("Cycle in graph")

    slayers = collections.OrderedDict((name, layers[name]) for name in snames)
    return slayers


[docs]class CudaConvnetNetwork(SequentialNetwork): def __init__(self, model, synapse=None, lif_type='lif', **kwargs): super(CudaConvnetNetwork, self).__init__(**kwargs) self.model = model self.synapse = synapse self.lif_type = lif_type # --- build model layers = dict(model['model_state']['layers']) # remove cost layer(s) and label inputs for name, layer in list(layers.items()): if layer['type'].startswith('cost'): layers.pop(name) layers.pop(layer['inputs'][0]) # first input is labels for layer in layers.values(): assert all(i in layers for i in layer.get('inputs', ())) # sort layers depths = layer_depths(layers) assert np.unique(list(depths.values())).size == len(depths) for name in sorted(layers, key=lambda n: depths[n]): self._add_layer(layers[name]) def _add_layer(self, layer): attrname = '_add_%s_layer' % layer['type'].replace('.', '_') if hasattr(self, attrname): return getattr(self, attrname)(layer) else: raise NotImplementedError( "Layer type %r not implemented" % layer['type']) def _get_inputs(self, layer): return [self.layers_by_name[i] for i in layer.get('inputs', [])] def _get_input(self, layer): assert len(layer.get('inputs', [])) == 1 return self._get_inputs(layer)[0] def _add_data_layer(self, layer): d = layer['outputs'] return self.add_data_layer(d, name=layer['name']) def _add_neuron_layer(self, layer): inputs = [self._get_input(layer)] neuron = layer['neuron'] ntype = neuron['type'] n = layer['outputs'] gain = 1. bias = 0. amplitude = 1. if ntype == 'ident': neuron_type = nengo.Direct() elif ntype == 'relu': neuron_type = nengo.RectifiedLinear() elif ntype == 'logistic': neuron_type = nengo.Sigmoid() elif ntype == 'softlif': from .neurons import SoftLIFRate tau_ref, tau_rc, alpha, amp, sigma, noise = [ neuron['params'][k] for k in ['t', 'r', 'a', 'm', 'g', 'n']] lif_type = self.lif_type.lower() if lif_type == 'lif': neuron_type = nengo.LIF(tau_rc=tau_rc, tau_ref=tau_ref) elif lif_type == 'lifrate': neuron_type = nengo.LIFRate(tau_rc=tau_rc, tau_ref=tau_ref) elif lif_type == 'softlifrate': neuron_type = SoftLIFRate( sigma=sigma, tau_rc=tau_rc, tau_ref=tau_ref) else: raise KeyError("Unrecognized LIF type %r" % self.lif_type) gain = alpha bias = 1. amplitude = amp else: raise NotImplementedError("Neuron type %r" % ntype) return self.add_neuron_layer( n, inputs=inputs, neuron_type=neuron_type, synapse=self.synapse, gain=gain, bias=bias, amplitude=amplitude, name=layer['name']) def _add_softmax_layer(self, layer): return None # non-neural, we can do without it # inputs = [self._get_input(layer)] # return self.add_softmax_layer( # d=layer['outputs'], inputs=inputs, name=layer['name']) def _add_dropout_layer(self, layer): inputs = [self._get_input(layer)] d = layer['outputs'] keep = layer['keep'] return self.add_dropout_layer( d, keep, inputs=inputs, name=layer['name']) def _add_dropout2_layer(self, layer): return self._add_dropout_layer(layer) def _add_fc_layer(self, layer): inputs = [self._get_input(layer)] weights = layer['weights'][0].T biases = layer['biases'].ravel() return self.add_full_layer( weights, biases, inputs=inputs, name=layer['name']) def _add_conv_layer(self, layer): inputs = [self._get_input(layer)] assert layer['sharedBiases'] nc = layer['channels'][0] nx = layer['imgSize'][0] ny = layer['modulesX'] nf = layer['filters'] s = layer['filterSize'][0] st = layer['stride'][0] p = -layer['padding'][0] filters = layer['weights'][0].reshape(nc, s, s, nf) filters = np.rollaxis(filters, axis=-1, start=0) biases = layer['biases'] layer = self.add_conv_layer( (nc, nx, nx), filters, biases, strides=st, padding=p, border='ceil', inputs=inputs, name=layer['name']) assert layer.node.output.shape_out == (nf, ny, ny) return layer def _add_local_layer(self, layer): inputs = [self._get_input(layer)] nc = layer['channels'][0] nx = layer['imgSize'][0] ny = layer['modulesX'] nf = layer['filters'] s = layer['filterSize'][0] st = layer['stride'][0] p = -layer['padding'][0] filters = layer['weights'][0].reshape(ny, ny, nc, s, s, nf) filters = np.rollaxis(filters, axis=-1, start=0) biases = layer['biases'].reshape(nf, ny, ny) return self.add_local_layer( (nc, nx, nx), filters, biases, strides=st, padding=p, inputs=inputs, name=layer['name']) def _add_pool_layer(self, layer): inputs = [self._get_input(layer)] assert layer['start'] == 0 nc = layer['channels'] nx = layer['imgSize'] s = layer['sizeX'] st = layer['stride'] kind = layer['pool'] return self.add_pool_layer( (nc, nx, nx), s, strides=st, kind=kind, mode='full', inputs=inputs, name=layer['name'])