能不能用 Tensorflow 严格复现 Alexnet 网络架构？答案

【问题标题】：Can we strictly reproduce Alexnet network architecture with Tensorflow?能不能用 Tensorflow 严格复现 Alexnet 网络架构？
【发布时间】：2021-02-19 14:46:26
【问题描述】：

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

从互联网上存在的numerous implementations 中，我找不到任何将模型分成两部分的模型，如论文和上图所述。据我了解，当时是因为 GPU 内存限制才这样做的。

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

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

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

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

【问题讨论】：

标签： python tensorflow distributed-computing imagenet

【解决方案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)

【讨论】：

您如何管理 GPU 分配？层计算 x1 需要在 GPU 1 上，层计算 x2 在 GPU 2 上。所有拆分层都一样