我们能用Tensorflow严格重现Alexnet网络架构吗？

Question

我想用Tensorflow 2严格重现Alexnet神经网络：

​

​

从互联网上存在的numerous implementations中，我找不到任何一个将模型一分为二的模型，就像上面的论文和图像所描述的那样。我知道这是因为当时GPU内存的限制。

但我想知道这样的实现在TensorFlow 2中是否可能，以及如何实现？如何将模型一分为二，分配给两个不同的GPU，同时仍然合并一些层(层2，5，6)的输出。

我敢打赌，Keras的级别太高了，无法处理这种复杂性，而TF的自动化分布式策略使得控制设备分配和同步变得困难。

我尝试了一下，将这些层分组为模型(例如，第1层和第2层)，并使用tf.device('/GPU:x')将这些组分配给它们各自的设备)，同时仍然连接输出。然后尝试在两个GPU上并行化和同步层...

但我无法靠近...因此，任何方法/建议都是值得赞赏的。

Answer 1

如果你使用函数式API，这并不是很复杂。在这里，我只演示如何将它一分为二，或者按您的喜好分成多个

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)