如何使用 TensorBoard 可视化具有自定义模型子类的 keras 模型？

Question

我有一个模型，它由几个继承自 tf.keras.Model 的子模型组成。这些子模型或多或少都是keras.Sequential 模型的简单集合，它们构成keras.layers，例如keras.layers.Conv2D、keras.layers.BatchNormalization 等。调用函数通过不同的顺序模型传递数据（有时添加额外的东西，例如顺序模型输出的原始输入，例如 ResidualBlock 子模型）。

我的主模型由子模型组成的原因是因为主模型很复杂，这样做可以让我轻松更改模型架构（例如子模型A的层数。另外，部分子模型实例化某些层（例如keras.layers.Reshape) in thecallfunction because the argument to configure theReshape`取决于调用函数的输入。

模型编译成功，我已经通过它传递了随机数据（尚未训练它），但我想将其可视化。

我尝试了以下操作

tensorboard = TensorBoard(log_dir='./logs/{}'.format(time()))
tensorboard.set_model(model)

但我收到警告：
WARNING:tensorflow:Model failed to serialize as JSON. Ignoring...

我也不能用
model.save('path_to_file.h5')保存它
因为我得到一个“NotImplemnedError”。

经过研究，我发现保存自定义模型的推荐方法是仅保存权重并仅加载权重。

如何使用 Tensorboard 可视化我的模型？我需要实现一个序列化器吗？有这方面的指南吗？

Answer 1

就使用 Keras API 的 Tensorboard 而言，您可能不太走运，因为它并非设计用于在具有自定义功能的嵌套模型上运行。好消息是，链接的源代码并不难理解，实际上对我来说比官方指南更直观 - 所以你应该能够编写自己的 tensorboard 类来满足你的需求。

嵌套回调形式的任何“变通办法”都将比自定义类更容易出错且更难管理 - 因此，虽然后者最初可能涉及更多工作，但从长远来看应该会有所回报。

最后，Tensorboard API 的可定制性有限 - 例如。无法选择要记录的特定层，或者要省略哪些指标。为此，我编写了自己的课程——见下文摘录；它不支持嵌套模型，但可以轻松扩展。

def run(self, log_num=None):
    tensors = (self.model.inputs + 
               self.model.targets + 
               self.model.sample_weights)
    assert len(tensors) == len(self.val_data)

    feed_dict = dict(zip(tensors, self.val_data))
    summary = self.sess.run([self.merged], feed_dict=feed_dict)

    log_num = log_num or self.log_num
    self.writer.add_summary(summary[0], log_num)
    self.log_num += 1

    if self.verbose:
        print("MyTensorBoard saved %s logs" % len(summary))

def _init_logger(self):
    for layer in self.model.layers:
        if any([(spec_name in layer.name) for spec_name in self.layer_names]):
            grads = self._get_grads(layer)
            if grads is not None:
                tf.summary.histogram(layer.name + '_grad', grads)
            if hasattr(layer, 'output'):
                self._log_outputs(layer)

            for weight in layer.weights:
                mapped_weight_name = weight.name.replace(':', '_')
                tf.summary.histogram(mapped_weight_name, weight)

                w_img = self._to_img_format(weight)
                if w_img is not None:
                    tf.summary.image(mapped_weight_name, w_img)
    self.merged = tf.summary.merge_all()
    self._init_writer()
    print("MyTensorBoard initialized")


def _init_writer(self):
    tb_num = 0
    while any([('TB_' + str(tb_num) in fname) for fname in 
               os.listdir(self.base_logdir)]):
        tb_num += 1
    self.logdir = os.path.join(self.base_logdir, 'TB_%s' % tb_num)
    os.mkdir(self.logdir)
    print("New TB logdir created at %s" % self.logdir)

    if self.write_graph:
        self.writer = tf.summary.FileWriter(self.logdir, self.sess.graph)
    else:
        self.writer = tf.summary.FileWriter(self.logdir)

def _get_grads(self, layer):
    for weight_tensor in layer.trainable_weights:
        grads = self.model.optimizer.get_gradients(
                     self.model.total_loss, weight_tensor)

        is_indexed_slices = lambda g: type(g).__name__ == 'IndexedSlices'
        return [grad.values if is_indexed_slices(grad) 
                            else grad for grad in grads]

def _log_outputs(self, layer):
    if isinstance(layer.output, list):
        for i, output in enumerate(layer.output):
            tf.summary.histogram('{}_out_{}'.format(layer.name, i), output)
    else:
        tf.summary.histogram('{}_out'.format(layer.name), layer.output)