Как распечатать все формы активации (более подробные, чем итоговые ()) для модели Tensorflow V2 keras?

Question

Я провел много времени с Tensorflow v.0 и v.1, и сейчас я пробую модель Tensorflow v.2 keras. model.summary() выглядело легко и удобно, но без деталей.

Вот игрушечный пример. Допустим, я определяю пользовательские слои и модели, как показано ниже (функциональный стиль API и стиль подкласса).

См. Ниже. Я хотел видеть примитивные слои внутри пользовательских слоев, но .summary() показывает только поверхностную информацию (только прямые дочерние слои).

Игрушечные пользовательские слои (слои - это просто определения игрушек):

import tensorflow as tf
from tensorflow import keras
from tensorflow.keras.layers import (Dense, Conv2D, BatchNormalization)

class LayerA(tf.keras.layers.Layer):
    def __init__(self, num_outputs, **kwargs):
        super(LayerA, self).__init__(**kwargs)
        self.num_outputs = num_outputs
        self.dense1 = Dense(64)
        self.dense2 = Dense(128)
        self.dense3 = Dense(num_outputs)
        self.bn = BatchNormalization(trainable=True)

    def call(self, inputs, training=True):
        x = self.dense1(inputs)
        x = self.bn(x, training=training)
        x = self.dense1(inputs)
        x = self.bn(x, training=training)
        x = self.dense1(inputs)
        x = self.bn(x, training=training)
        return x

class LayerB(tf.keras.layers.Layer):
    def __init__(self, num_outputs, **kwargs):
        super(LayerB, self).__init__(**kwargs)
        self.num_outputs = num_outputs
        self.dense = Dense(64)
        self.bn = BatchNormalization(trainable=True)

    def call(self, inputs, training=True):
        x = self.dense(inputs)
        x = self.bn(x, training=training)
        return x

Определение модели с функциональным API:

inputs = tf.keras.Input(shape=(28), name='input')
x = LayerA(7, name='layer_a')(inputs)
x = LayerB(13, name='layer_b')(x)
x = tf.reduce_max(x, 1)
model_func = keras.Model(inputs=inputs, outputs=x, name='model')
model_func.summary()

# Results:
# Layer (type)                 Output Shape              # Param #   
# =================================================================
# input (InputLayer)           [(None, 28)]              0         
# _________________________________________________________________
# layer_a (LayerA)             (None, 64)                2112      
# layer_b (LayerB)             (None, 64)                4416      
# tf_op_layer_Max_1 (TensorFlo [(None,)]                 0         
# =================================================================
# Total params: 6,528
# Trainable params: 6,272
# Non-trainable params: 256

Подклассы определения модели:

class ModelA(tf.keras.Model):

    def __init__(self):
        super(ModelA, self).__init__()
        self.block_1 = LayerA(7, name='layer_a')
        self.block_2 = LayerB(13, name='layer_b')

    def call(self, inputs):
        x = self.block_1(inputs)
        x = self.block_2(x)
        x = tf.reduce_max(x, 1)
        return x

model_subclass = ModelA()
y = model_subclass(inputs)
model_subclass.summary()

### Result:
# Layer (type)                 Output Shape              Param #   
# =================================================================
# layer_a (LayerA)             (None, 64)                2112      
# layer_b (LayerB)             (None, 64)                4416      
# =================================================================
# Total params: 6,528
# Trainable params: 6,016
# Non-trainable params: 512

Как можноЯ печатаю все формы активации слоев Conv и Dense в модели? Например,

layer_a/dense_1                          (None, ...)
layer_a/dense_2                          (None, ...)
layer_b/dense_1                          (None, ...)
layer_b/maybe-even-deeper-layer/conv2d_1 (None, ...)
... etc ...

В Tensorflow v.0 или v.1 я бы сделал что-то вроде:

for n in tf.get_default_graph().as_graph_def().node:
    print(n.name, n.shape)

Есть ли способ напечатать больше деталей, когда у меня есть керасмодель?

Answer 1

Сводка не будет делать это автоматически, поэтому вы должны адаптироваться.

Вы можете, например, создать периодическую сводку:

def full_summary(layer):

    #check if this layer has layers
    if hasattr(layer, 'layers'):
        print('summary for ' + layer.name)
        layer.summary()
        print('\n\n')

        for l in layer.layers:
            full_summary(l)

Использовать ее как:

full_summary(my_model)