如何查看HuggingFace TFRobertaModel的详细网络层摘要信息

问题复现 #

初始代码如下：

from transformers import RobertaTokenizer, TFRobertaModel
import tensorflow as tf

tokenizer = RobertaTokenizer.from_pretrained("roberta-base")
model = TFRobertaModel.from_pretrained("roberta-base")

直接执行model.summary()时，仅会展示顶层封装的模型结构，无法查看内部嵌入层、Transformer编码层、池化层等子层的输出形状、参数量信息，默认输出如下：

Model: "tf_roberta_model_2"
_________________________________________________________________
Layer (type)                 Output Shape              Param #   
=================================================================
roberta (TFRobertaMainLayer) multiple                  124645632 
=================================================================
Total params: 124,645,632
Trainable params: 124,645,632
Non-trainable params: 0
_________________________________________________________________

可用解决方法 #

• 直接调用内部主层的摘要方法
  HuggingFace的TensorFlow预训练模型会把所有实际计算逻辑封装在模型实例的同名主层属性下，RoBERTa对应的内部主层为model.roberta。调用该层的summary方法时传入expand_nested=True参数，即可展开所有嵌套的子层结构：

  # 如需显示明确的输出形状，先传入样本完成模型构建
  sample_input = tokenizer("test text", return_tensors="tf")
  _ = model(sample_input)
  
  # 打印完整层级摘要
  model.roberta.summary(expand_nested=True, show_trainable=True)
  

  该方法对所有HuggingFace的TF结构模型通用，比如TFBertModel对应内部主层为model.bert，替换对应属性名即可使用。
• 递归遍历提取所有层信息
  如果需要单独获取某一层的属性，可以通过递归遍历layers属性实现：

  def list_all_layers(layer, indent=0):
      print("  "*indent + f"- {layer.name} | 层类型: {layer.__class__.__name__} | 参数量: {layer.count_params()}")
      # 存在嵌套子层则递归遍历
      if hasattr(layer, "layers"):
          for sub_layer in layer.layers:
              list_all_layers(sub_layer, indent+1)
  
  list_all_layers(model)
  

• 可视化导出层级结构
  如需可视化查看结构，可调用Keras的绘图接口，指定展开嵌套层、显示形状参数：

  tf.keras.utils.plot_model(
      model.roberta,
      to_file="roberta_layer_structure.png",
      show_shapes=True,
      show_layer_names=True,
      expand_nested=True,
      dpi=96
  )
  

内容的提问来源于stack exchange，提问作者Murali