Hello everyone!
I am currently trying to convert a CSPNet style network from Keras to native Nengo. This means at some point I need to split my computation graph and slice the tensors like in this Keras code:
x1 = Conv2D(...)(inp)
# Slice tensor in half
x2 = x1[...,:x1.shape[-1]//2]
x2 = Conv2D(...)(x2)
...
# Concat the two tensors and fuse the graph
x1 = Concatenate(axis=-1)([x1, x2])
x1 = Conv2D(...)(x1)
...
model = Model(inp, x1)
The problem with this is that Keras internally uses the SlicingOpLambda to wrap the slice operation, which is (obviously) not implemented in the converter.py file. So naturally I looked into implementing my own custom LayerConverter class for converting this operation to Nengo.
This should be pretty simple since the only thing to do is to create two nodes with the second one being half the size of the preceding one and then connecting half of the preceding neurons. The code for this looks as follows:
@nengo_dl.Converter.register(SlicingOpLambda)
class ConvertSlicingOpLambda(nengo_dl.converter.LayerConverter):
def convert(self, node_id):
output = self.add_nengo_obj(node_id)
self.add_connection(node_id=node_id, obj=output, pre_slice=slice(output.size_out))
return output
Note the pre_slice argument for the add_connection method. It is not there in the original method, with the reason being that the normal add_connection method doesn’t support connecting only part of the preceding node to the object given with obj. That is why I modified the original method like below:
def add_connection(self, node_id, obj, pre_slice=None, input_idx=0, trainable=False, **kwargs):
...
if pre_slice is not None:
# Only connect the part given by the pre_slice argument
conn = nengo.Connection(
pre[pre_slice],
obj,
**kwargs,
)
else:
# Just like before
...
With this “hack” my model does not perform the same in native Nengo as in Keras though, meaning something did not quite go right.
Is there any other method of connecting only part of a previous node with the one created in the convert method of a LayerConverter?
Thanks in advance, any help is much appreciated!
Cheers.
Just in case it is helpful, here is a full minimum example for showing the issue. Use the use_slice variable to switch between using or not using the custom operation.
import tensorflow as tf
import nengo_dl
import numpy as np
from tensorflow.python.keras.layers.core import SlicingOpLambda
physical_devices = tf.config.experimental.list_physical_devices("GPU")
if len(physical_devices) > 0:
tf.config.experimental.set_memory_growth(physical_devices[0], True)
print("Called tf.config.experimental.set_memory_growth(GPU0, True)")
def evaluate_accuracy(predictions, ground_truth):
prediction_values = np.argmax(tf.nn.softmax(predictions), axis=-1)
ground_truth_values = np.argmax(ground_truth[:len(predictions)], -1)
return (prediction_values == ground_truth_values).mean()
@nengo_dl.Converter.register(SlicingOpLambda)
class ConvertSlicingOpLambda(nengo_dl.converter.LayerConverter):
def convert(self, node_id):
output = self.add_nengo_obj(node_id)
self.add_connection(node_id=node_id, obj=output, pre_slice=slice(output.size_out))
return output
# Parameter for comparing
use_slice = True
# Load MNIST dataset
dataset = tf.keras.datasets.mnist.load_data(path="mnist.npz")
inp = tf.keras.Input((28, 28, 1))
x1 = tf.keras.layers.Conv2D(16, 3)(inp)
if use_slice:
x2 = x1[...,:x1.shape[-1]//2]
else:
x2 = x1
x2 = tf.keras.layers.Conv2D(16, 3, padding="same")(x2)
x1 = tf.keras.layers.Concatenate(axis=-1)([x1, x2])
x1 = tf.keras.layers.Conv2D(16, 1)(x1)
x1 = tf.keras.layers.Flatten()(x1)
x1 = tf.keras.layers.Dense(10)(x1)
model = tf.keras.Model(inp, x1)
# Train with keras
model.compile(
loss=tf.losses.CategoricalCrossentropy(from_logits=True),
optimizer='adam', metrics=['accuracy'])
model.fit(dataset[0][0], tf.keras.utils.to_categorical(
dataset[0][1], num_classes=10), batch_size=32, epochs=1)
print("Accuracy in second simulator: {}".format(
model.evaluate(
dataset[1][0][:100],
tf.keras.utils.to_categorical(dataset[1][1][:100], num_classes=10)
)
))
# Reshape test data
test_data = dataset[1][0].reshape((dataset[1][0].shape[0], 1, -1))
test_target = np.expand_dims(
tf.keras.utils.to_categorical(dataset[1][1], num_classes=10), 1)
# Create converter
converter = nengo_dl.Converter(model, allow_fallback=False)
nengo_inp = converter.inputs[inp]
nengo_outp = converter.outputs[x1]
# Simulator for testing with Nengo
with nengo_dl.Simulator(converter.net, minibatch_size=1) as sim:
result = sim.predict({nengo_inp: test_data[:100]})[nengo_outp]
print("Accuracy in second simulator: {}".format(
evaluate_accuracy(result, test_target[:100])
))
