HuggingFace TFRobertaModel详细摘要

Question

from transformers import RobertaTokenizer, TFRobertaModel
import tensorflow as tf

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

我需要这个HuggingFace TFRobertaModel()的详细图层摘要，这样我就可以可视化形状、图层并在需要时进行定制。然而，当我这样做：model.summary()时，它只是在一个层中显示所有内容。我试着挖掘它的不同属性，但无法获得详细的图层摘要。是否可以这样做呢？

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论坛上还没有得到答复。

Answer 1

不完全是一个模型摘要，但您可以打印如下所示的层：

from transformers import RobertaTokenizer, TFRobertaModel
import tensorflow as tf

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

def print_layers(l, model):
  for idx, s in enumerate(l.submodules):
    if s.submodules:
      print_layers(s, model)
    print(s)

TFRobertaMainLayer = model.layers[0]   
print_layers(TFRobertaMainLayer, model)

您还可以使用s.weights来获取每个层的权重。