python - 我们可以用 TensorFlow 严格复现 Alexnet 网络架构吗？

Question

我想严格复制 Alexnet使用 Tensorflow 2 的神经网络:

来自numerous implementations存在于互联网上，我找不到任何模型如论文和上图所述一分为二。据我所知，当时是因为GPU内存限制才做的。

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

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

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

但我无法接近....所以，任何方法/建议都值得赞赏

Answer 1

如果你使用Functional 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)