Похоже, что снижение производительности во фрагменте кода, непосредственно создающем экземпляр MainLayer, происходит из-за того, что предварительно обученные веса не загружаются. Вы можете загрузить веса одним из следующих способов:

1. Вызов TFBertModel.from_pretrained и захват MainLayer из загруженного TFBertModel
2. Создание MainLayer напрямую, затем загрузка весов аналогично from_pretrained

Почему это происходит

Когда вы вызываете TFBertModel.from_pretrained, он использует функцию TFPreTrainedModel.from_pretrained (через наследование), которая обрабатывает несколько вещей, включая загрузку, кэширование и загрузку весов модели.

class TFPreTrainedModel(tf.keras.Model, TFModelUtilsMixin, TFGenerationMixin):
    ...
    @classmethod
    def from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs):
        ...
        # Load model
        if pretrained_model_name_or_path is not None:
            if os.path.isfile(os.path.join(pretrained_model_name_or_path, TF2_WEIGHTS_NAME)):
            # Load from a TF 2.0 checkpoint
            archive_file = os.path.join(pretrained_model_name_or_path, TF2_WEIGHTS_NAME)
            ...
            resolved_archive_file = cached_path(
                    archive_file,
                    cache_dir=cache_dir,
                    force_download=force_download,
                    proxies=proxies,
                    resume_download=resume_download,
                    local_files_only=local_files_only,
            )
            ...
            model.load_weights(resolved_archive_file, by_name=True)

(Если вы читали реальный код, многое было ... вычеркнуто выше).

Однако, когда вы создаете экземпляр TFBertMainLayer, он не делает ничего из этого настроить работу.

@keras_serializable
class TFBertMainLayer(tf.keras.layers.Layer):
    config_class = BertConfig

    def __init__(self, config, **kwargs):
        super().__init__(**kwargs)
        self.num_hidden_layers = config.num_hidden_layers
        self.initializer_range = config.initializer_range
        self.output_attentions = config.output_attentions
        self.output_hidden_states = config.output_hidden_states
        self.return_dict = config.use_return_dict
        self.embeddings = TFBertEmbeddings(config, name="embeddings")
        self.encoder = TFBertEncoder(config, name="encoder")
        self.pooler = TFBertPooler(config, name="pooler")
   
   ... rest of the class

По сути, вам нужно убедиться, что эти веса загружаются.

Решения

(1) Использование TFAutoModel.from_pretrained

Вы можете положиться на transforms.TFAutoModel.from_pretrained для загрузки модели, а затем просто взять поле MainLayer из определенного подкласса TFPreTrainedModel. Например, если вы хотите получить доступ к основному слою дистиллятора, это будет выглядеть так:

    model = transformers.TFAutoModel.from_pretrained(`distilbert-base-uncased`)
    assert isinstance(model, TFDistilBertModel)
    main_layer = transformer_model.distilbert

Вы можете видеть в modeling_tf_distilbert.html, что MainLayer является полем модели . Это меньше кода и меньше дублирования, но есть несколько недостатков. Менее просто изменить предварительно обученную модель, которую вы собираетесь использовать, потому что теперь вы зависите от имени поля, если вы измените тип модели, вам придется изменить имя поля (например, в TFAlbertModel MainLayer поле называется albert). Кроме того, кажется, что это не предназначенный способ использования Huggingface, так что это может измениться у вас под носом, и ваш код может сломаться из-за обновлений Huggingface.

class TFDistilBertModel(TFDistilBertPreTrainedModel):
    def __init__(self, config, *inputs, **kwargs):
        super().__init__(config, *inputs, **kwargs)
        self.distilbert = TFDistilBertMainLayer(config, name="distilbert")  # Embeddings

[DOCS]    @add_start_docstrings_to_callable(DISTILBERT_INPUTS_DOCSTRING)
    @add_code_sample_docstrings(
        tokenizer_class=_TOKENIZER_FOR_DOC,
        checkpoint="distilbert-base-uncased",
        output_type=TFBaseModelOutput,
        config_class=_CONFIG_FOR_DOC,
    )
    def call(self, inputs, **kwargs):
        outputs = self.distilbert(inputs, **kwargs)
        return outputs

(2) Повторная реализация логики загрузки веса из from_pretrained

Вы можете сделать это, скопировав/вставив те части from_pretrained, которые имеют отношение к загрузке веса. Это также имеет некоторые серьезные недостатки, вы будете дублировать логику, которая может не синхронизироваться с библиотеками Huggingface. Хотя вы, вероятно, могли бы написать его более гибким и устойчивым к изменениям имени базовой модели.

Вывод

В идеале это то, что будет исправлено внутри командой Huggingface, либо путем предоставления стандартной функции для создания MainLayer, переноса логики загрузки веса в собственную функцию, которую можно вызвать, либо путем поддержки сериализации в классе модели.