由于 Python 操作而弄乱了 TensorBoard 图

Question

这个问题是关于 TensorFlow（和 TensorBoard）版本 2.2rc3，但我在 2.1 中遇到了同样的问题。

考虑以下奇怪的代码：

from datetime import datetime

import tensorflow as tf
from tensorflow import keras

inputs = keras.layers.Input(shape=(784, ))    

x1 = keras.layers.Dense(32, activation='relu', name='Model/Block1/relu')(inputs)
x1 = keras.layers.Dropout(0.2, name='Model/Block1/dropout')(x1)
x1 = keras.layers.Dense(10, activation='softmax', name='Model/Block1/softmax')(x1)

x2 = keras.layers.Dense(32, activation='relu', name='Model/Block2/relu')(inputs)
x2 = keras.layers.Dropout(0.2, name='Model/Block2/dropout')(x2)
x2 = keras.layers.Dense(10, activation='softmax', name='Model/Block2/softmax')(x2)

x3 = keras.layers.Dense(32, activation='relu', name='Model/Block3/relu')(inputs)
x3 = keras.layers.Dropout(0.2, name='Model/Block3/dropout')(x3)
x3 = keras.layers.Dense(10, activation='softmax', name='Model/Block3/softmax')(x3)

x4 = keras.layers.Dense(32, activation='relu', name='Model/Block4/relu')(inputs)
x4 = keras.layers.Dropout(0.2, name='Model/Block4/dropout')(x4)
x4 = keras.layers.Dense(10, activation='softmax', name='Model/Block4/softmax')(x4)

outputs = x1 + x2 + x3 + x4

model = tf.keras.Model(inputs=inputs, outputs=outputs)
model.summary()

(x_train, y_train), (x_test, y_test) = keras.datasets.mnist.load_data()
x_train = x_train.reshape(60000, 784).astype('float32') / 255
x_test = x_test.reshape(10000, 784).astype('float32') / 255

model.compile(loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
              optimizer=keras.optimizers.RMSprop(),
              metrics=['accuracy'])

logdir = "logs/" + datetime.now().strftime("%Y%m%d-%H%M%S")
tensorboard_callback = keras.callbacks.TensorBoard(log_dir=logdir)
model.fit(x_train, y_train,
          batch_size=64,
          epochs=5,
          validation_split=0.2,
          callbacks=[tensorboard_callback])

运行它并查看在 TensorBoard 中创建的图表时 you will see the following。

可以看出，加法操作真的很丑。

换行时

outputs = x1 + x2 + x3 + x4

用线条：

outputs = keras.layers.add([x1, x2], name='Model/add/add1')
outputs = keras.layers.add([outputs, x3], name='Model/add/add2')
outputs = keras.layers.add([outputs, x4], name='Model/add/add3')

a much nicer graph is created by TensorBoard（在第二个屏幕截图中，详细显示了模型以及其中一个内部块）。

模型的两种表示的区别在于，在第二种表示中，我们可以命名加法运算并将它们分组。

我找不到任何方法来命名这些操作，除非使用keras.layers.add()。在这个模型中，问题看起来并不那么严重，因为模型很简单，并且很容易将+ 替换为keras.layers.add()。然而，在更复杂的模型中，它可能会成为一种真正的痛苦。例如，t[:, start:end] 等操作应转换为对tf.strided_slice() 的复杂调用。所以我的模型表示非常混乱，有很多神秘的聚集、跨步和连接操作。

我想知道是否有办法包装/分组此类操作以在 TensorBoard 中提供更好的图形。

Answer 1

outputs = keras.layers.Add()([x1, x2, x3, x4])

Answer 2

根据 Marco Cerliani 的提示，Lambda 层在这里确实非常有用。所以下面的代码将很好地分组+：

outputs = keras.layers.Lambda(lambda x: x[0] + x[1], name='Model/add/add1')([x1, x2])
outputs = keras.layers.Lambda(lambda x: x[0] + x[1], name='Model/add/add2')([outputs, x2])
outputs = keras.layers.Lambda(lambda x: x[0] + x[1], name='Model/add/add3')([outputs, x2])

或者如果需要包装跨步，以下代码将很好地分组t[]：

x1 = keras.layers.Lambda(lambda x: x[:, 0:5], name='Model/stride_concat/stride1')(x1) # instead of x1 = x1[:, 0:5]
x2 = keras.layers.Lambda(lambda x: x[:, 5:10], name='Model/stride_concat/stride2')(x2) # instead of x2 = x2[:, 5:10]
outputs = keras.layers.concatenate([x1, x2], name='Model/stride_concat/concat')

这回答了所提出的问题。但实际上，另一个问题中描述了一个未解决的问题：'TensorFlowOpLayer messes up the TensorBoard graphs'