FlagAI框架高级使用指南：自定义模型与Tokenizer选择-优快云博客

本文链接：https://blog.youkuaiyun.com/gitblog_00025/article/details/148552146

FlagAI框架高级使用指南：自定义模型与Tokenizer选择

FlagAI FlagAI (Fast LArge-scale General AI models) is a fast, easy-to-use and extensible toolkit for large-scale model. 项目地址: https://gitcode.com/gh_mirrors/fl/FlagAI

自定义模型开发

在FlagAI框架中，开发者可以基于现有模型进行扩展或完全自定义新模型。这一功能为研究人员和工程师提供了极大的灵活性，使他们能够针对特定任务优化模型架构。

自定义模型开发规范

开发自定义模型时，需要遵循以下核心规范：

继承BaseModel基类
所有自定义模型必须继承自BaseModel，这确保了模型能够支持框架提供的标准接口，包括预训练参数加载(from_pretrain)和基于配置文件初始化(init_from_json)等功能。
初始化函数要求
__init__()函数的第一个参数必须是config，该参数对应模型配置文件(config.json)中的参数。此外，开发者可以自由添加其他任务特定参数。
权重加载函数
必须实现load_weights()函数，负责加载预训练权重。这个函数应当处理权重文件的路径解析和参数加载逻辑。
前向传播输出格式
forward()函数必须返回一个字典，其中必须包含logits键。如果输入中包含标签数据(labels)，则还需要返回loss值。

实战示例：GLM序列分类模型

让我们通过一个GLM模型完成序列分类任务的例子，深入理解自定义模型的实现：

from flagai.model.base_model import BaseModel
from flagai.model.glm_model import GLMModel
import torch

class GLMForSequenceClassification(BaseModel):
    def __init__(self, config, hidden_dropout=0.1, pool_token='cls', **kwargs):
        super().__init__(config, **kwargs)
        self.config = config
        self.pool_token = pool_token
        self.model = GLMModel(config)
        self.model.output_predict = False
        self.num_class = config['class_num']
        
        # 构建分类头
        hidden_size = self.model.hidden_size
        self.pool_layer = torch.nn.Linear(hidden_size, hidden_size)
        self.multichoice_dropout = torch.nn.Dropout(hidden_dropout)
        self.multichoice_head = torch.nn.Linear(hidden_size, self.num_class)

在这个初始化函数中，我们除了接收必要的config参数外，还定义了两个重要参数：

hidden_dropout: 控制分类头前的dropout率
pool_token: 指定如何从序列中提取特征（'cls'表示使用[CLS]标记，'start'表示使用起始标记等）

前向传播函数的实现需要考虑多种输入情况：

def forward(self, input_ids=None, position_ids=None, attention_mask=None, **kwargs):
    # 处理多选任务输入
    if len(input_ids.shape) == 3:
        batch_size, num_choices = input_ids.shape[:2]
        input_ids = input_ids.reshape(-1, input_ids.size(-1))
        attention_mask = attention_mask.reshape(-1, *attention_mask.size()[2:])
        position_ids = position_ids.reshape(-1, *position_ids.size()[2:])
    
    # 获取GLM模型输出
    model_out = self.model(input_ids, position_ids, attention_mask)
    outputs, mems = model_out['logits'], model_out['hidden_states']
    
    # 根据pool_token策略提取特征
    if self.pool_token == 'start':
        output = outputs[torch.arange(outputs.size(0), attention_mask]
    elif self.pool_token == 'pad':
        output = outputs[torch.arange(outputs.size(0), attention_mask - 1]
    elif self.pool_token == 'cls':
        output = outputs[:, 0]
    
    # 通过分类头得到最终logits
    output = torch.tanh(self.pool_layer(output))
    multichoice_output = self.multichoice_dropout(output)
    logits = self.multichoice_head(multichoice_output)
    
    # 返回结果
    if 'labels' not in kwargs:
        return {'logits': logits, 'hidden_states': mems}
    else:
        labels = kwargs['labels']
        # 计算损失
        if logits.size(1) == 1:
            loss = F.binary_cross_entropy_with_logits(logits.float(), labels.float())
        else:
            loss = F.cross_entropy(logits.float(), labels.long())
        return {"loss": loss, 'logits': logits, 'hidden_states': mems}

模型使用方式

完成自定义模型开发后，可以通过框架提供的便捷接口加载和使用：

model_dir = "./state_dict/GLM_sequence_classification/"
model = GLMForSequenceClassification.from_pretrain(
    model_dir, 
    hidden_dropout=0.1,
    pool_token="cls"
)

这种方式既保持了预训练模型的能力，又可以根据具体任务需求进行灵活调整。