Can we strictly reproduce Alexnet network architecture with Tensorflow?

user_3010217 0
Viewed 84

I want to strictly reproduce the Alexnet neural network with Tensorflow 2:

enter image description here

From the numerous implementations that exist on the internet, I could not find any having the model split in two as described by the paper and the image above. I understand that it was done because of the GPU memory limitation at that time.

But I'm wondering if such implementation is possible with TensorFlow 2 and how? How a model could be split in two, assigned to two different GPUs while still merging the output of some layers (layers 2, 5, 6).

I bet that Keras is too high level to handle this complexity, and the automated distributed strategy of TF makes it hard to control the device assignments and synchronization.

I gave a try by grouping the layers into models (e.g. layer 1 and 2) and assign those groups to their respective device using tf.device('/GPU:x')) while still concatenating the output. And then tried to parallelize and sync the layers across the two GPUs...

But I couldn't get any close.... So, any approach/suggestion is appreciated

imagenet tensorflow distributed-computing python
1 Answers

It is not very complicated if you use Functional API. Here I only demonstrate how to split it in two, or in many as you like

inp = tf.keras.Input(input_shape=shape)
x1 = Layers.Conv2D(filters_1, kernel_size_1)(inp)
x2 = Layers.Conv2D(filters_2, kernel_size_2)(inp)
...
# Concatenate outputs of previous layers, be careful of dimension matching.
x = Layers.Concatenate()([x1, x2])
x = Layers.Dense(512)(x)
x = Layers.Dense(256)(x)
...
output = Layers.Dense(classes, activation="softmax")(x)
model = tf.keras.Model(inputs = inp, outputs = output)
user_112352050
Related