自定义模型继承 Hugging 的 PreTrainedModel

悄悄地努力2024-04-12 21:50

问题背景

使用 HuggingFace 的 Trainer 进行训练，自定义模型继承 PreTrainedModel，希望能够通过 from_pretrained 容易地实现：从上次保存的 checkpoint 处继续训练。

问题描述

复制代码

Traceback (most recent call last):
  File "D:\PycharmProjects\gaze\SPIGA\train\train.py", line 60, in <module>
    .from_pretrained('checkpoint-1175')
  File "D:\Anaconda3\envs\transformers\lib\site-packages\transformers\modeling_utils.py", line 2535, in from_pretrained
    config, model_kwargs = cls.config_class.from_pretrained(
AttributeError: 'NoneType' object has no attribute 'from_pretrained'

问题分析

问题定位为：cls.config_class 为 None，结果导致 'NoneType' object has no attribute 'from_pretrained'。

问题解决

自定义模型时，声明 config_class

python 复制代码

from transformers.configuration_utils import PretrainedConfig
from transformers.modeling)utils import PreTrainedModel

class CustomModel(PreTrainedModel):
	# 注意 1：在这里声明 config_class
    config_class = PretrainedConfig
    # 注意 2：重写的 __init__ 方法需要有 config 参数，不然会报错
    def __init__(self, config=PretrainedConfig(), num_landmarks=98, num_edges=15, steps=3, **kwargs):
        super().__init__(config=config)
        # .......