使用AI Gallery SDK构建自定义模型
AI Gallery的Transformers库支持部分开源的模型结构框架,并对昇腾系列显卡进行了训练/推理性能优化,可以做到开箱即用。如果你有自己从头进行预训练的模型,AI Gallery也支持使用SDK构建自定义模型接入AI Gallery。
Transformers库介绍
AI Gallery使用的Transformers机器学习库是一个开源的基于Transformer模型结构提供的预训练语言库。Transformers库注重易用性,屏蔽了大量AI模型开发使用过程中的技术细节,并制定了统一合理的规范。使用者可以便捷地使用、下载模型。同时支持用户上传自己的预训练模型到在线模型资产仓库中,并发布上架给其他用户使用。AI Gallery在原有Transformers库的基础上,融入了对于昇腾硬件的适配与支持。对AI有使用诉求的企业、NLP领域开发者,可以借助这个库,便捷地使用昇腾算力进行自然语言理解(NLU)和自然语言生成(NLG)任务的SOTA模型开发与应用。
支持的模型结构框架
AI Gallery的Transformers库支持的开源模型结构框架如表1所示。
核心基础类介绍
使用AI Gallery SDK构建自定义模型,需要了解2个核心基础类“PretrainedModel”和“PretrainedConfig”之间的交互。
- “PretrainedConfig”:预训练模型的配置基类
提供模型配置的通用属性和两个主要方法,用于序列化和反序列化配置文件。
PretrainedConfig.from_pretrained(dir) # 从目录中加载序列化对象(本地或者是url),配置文件为dir/config.json PretrainedConfig.save_pretrained(dir) # 将配置实例序列化到dir/config.json
- “PretrainedModel”:预训练模型的基类
包含一个配置实例“config”,提供两个主要方法,用来加载和保存预训练模型。
# 1. 调用 init_weights() 来初始化所有模型权重 # 2. 从目录中(本地或者是url)中导入序列化的模型 # 3. 使用导入的模型权重覆盖所有初始化的权重 # 4. 调用 PretrainedConfig.from_pretrained(dir)来将配置设置到self.config中 PretrainedModel.from_pretrained(dir) # 将模型实例序列化到 dir/pytorch_model.bin 中 PretrainedModel.save_pretrained(dir) # 给定input_ids,生成 output_ids,在循环中调用 PretrainedModel.forward() 来做前向推理 PretrainedModel.generate()
操作步骤
本文使用NewBert模型介绍构建自定义模型的流程。
- 安装AI Gallery SDK。
通过pip在本地或云上开发环境安装AI Gallery SDK(galleryformers)。
pip install galleryformers
建议在虚拟环境(Python 3.8+)中安装AI Gallery SDK,以便管理不同的项目,避免依赖项之间产生兼容性问题。
- 构建自定义模型。
- 编写自定义配置类。
模型的configuration包含了构建模型所需的所有信息的对象,需要尽可能完整。
from galleryformers import PretrainedConfig from typing import List class NewBertConfig(PretrainedConfig): model_type = "bert" def __init__( self, vocab_size=30522, hidden_size=768, num_hidden_layers=12, num_attention_heads=12, intermediate_size=3072, hidden_act="gelu", hidden_dropout_prob=0.1, attention_probs_dropout_prob=0.1, max_position_embeddings=512, type_vocab_size=2, initializer_range=0.02, layer_norm_eps=1e-12, pad_token_id=0, position_embedding_type="absolute", use_cache=True, classifier_dropout=None, **kwargs, ): super().__init__(pad_token_id=pad_token_id, **kwargs) self.vocab_size = vocab_size self.hidden_size = hidden_size self.num_hidden_layers = num_hidden_layers self.num_attention_heads = num_attention_heads self.hidden_act = hidden_act self.intermediate_size = intermediate_size self.hidden_dropout_prob = hidden_dropout_prob self.attention_probs_dropout_prob = attention_probs_dropout_prob self.max_position_embeddings = max_position_embeddings self.type_vocab_size = type_vocab_size self.initializer_range = initializer_range self.layer_norm_eps = layer_norm_eps self.position_embedding_type = position_embedding_type self.use_cache = use_cache self.classifier_dropout = classifier_dropout
- 自定义配置类必须继承自“PretrainedConfig”。
- 自定义配置类的“__init__”必须接受任何“kwargs”,这些“kwargs”需要传递给“__init__”。
- 完成自定义配置类的编写后,可以使用该类创建配置实例。
newbert1_config = NewBertConfig(num_hidden_layers=6, num_attention_heads=10, use_cache=False) newbert1_config.save_pretrained("mynewbert")
这一步会在本地名为mynewbert的文件夹中保存一个名为config.json的文件。
该配置实例同样可以通过调用from_pretrained方法加载。
newbert1_config.from_pretrained("mynewbert")
- 编写完配置部分,开始编写自定义模型。
下面展示了3种模型基类的代码示例,为了确保示例不过于复杂,本文对部分代码片段进行了省略展示。
- 预训练模型基类NewBertPreTrainedModel
from galleryformers import PreTrainedModel from .configuration_newbert import NewBertConfig class NewBertPreTrainedModel(PreTrainedModel): config_class = NewBertConfig load_tf_weights = load_tf_weights_in_bert base_model_prefix = "bert" supports_gradient_checkpointing = True def _init_weights(self, module): """Initialize the weights""" if isinstance(module, nn.Linear): module.weight.data.normal_(mean=0.0, std=self.config.initializer_range) if module.bias is not None: module.bias.data.zero_() elif isinstance(module, nn.Embedding): module.weight.data.normal_(mean=0.0, std=self.config.initializer_range) if module.padding_idx is not None: module.weight.data[module.padding_idx].zero_() elif isinstance(module, nn.LayerNorm): module.bias.data.zero_() module.weight.data.fill_(1.0)
- 基础模型类NewBertModel:该类继承自NewBertPreTrainedModel。
class NewBertModel(NewBertPreTrainedModel): def __init__(self, config, add_pooling_layer=True): super().__init__(config) self.config = config self.embeddings = BertEmbeddings(config) self.encoder = BertEncoder(config) self.pooler = BertPooler(config) if add_pooling_layer else None # Initialize weights and apply final processing self.post_init() def get_input_embeddings(self): return self.embeddings.word_embeddings def set_input_embeddings(self, value): self.embeddings.word_embeddings = value def _prune_heads(self, heads_to_prune): for layer, heads in heads_to_prune.items(): self.encoder.layer[layer].attention.prune_heads(heads) def forward( self, input_ids: Optional[torch.Tensor] = None, attention_mask: Optional[torch.Tensor] = None, token_type_ids: Optional[torch.Tensor] = None, position_ids: Optional[torch.Tensor] = None, head_mask: Optional[torch.Tensor] = None, inputs_embeds: Optional[torch.Tensor] = None, encoder_hidden_states: Optional[torch.Tensor] = None, encoder_attention_mask: Optional[torch.Tensor] = None, past_key_values: Optional[List[torch.FloatTensor]] = None, use_cache: Optional[bool] = None, output_attentions: Optional[bool] = None, output_hidden_states: Optional[bool] = None, return_dict: Optional[bool] = None, ...)
所有的模型都需要通过“forward”方法来实现自己的推理逻辑,这个方法会在执行“model(input_ids)”的时候进行调用
- 模型基类NewBertForXXX:该类承自NewBertPreTrainedModel。
该类可用于执行AI Gallery工具链服务,此处以文本问答(Question Answering)的任务类型为例:
class NewBertForQuestionAnswering(NewBertPreTrainedModel): def __init__(self, config): super().__init__(config) self.num_labels = config.num_labels self.bert = BertModel(config, add_pooling_layer=False) self.qa_outputs = nn.Linear(config.hidden_size, config.num_labels) # Initialize weights and apply final processing self.post_init() def forward( self, input_ids: Optional[torch.Tensor] = None, attention_mask: Optional[torch.Tensor] = None, token_type_ids: Optional[torch.Tensor] = None, position_ids: Optional[torch.Tensor] = None, head_mask: Optional[torch.Tensor] = None, inputs_embeds: Optional[torch.Tensor] = None, start_positions: Optional[torch.Tensor] = None, end_positions: Optional[torch.Tensor] = None, output_attentions: Optional[bool] = None, output_hidden_states: Optional[bool] = None, return_dict: Optional[bool] = None, ) return_dict = return_dict if return_dict is not None else self.config.use_return_dict outputs = self.bert( input_ids, attention_mask=attention_mask, token_type_ids=token_type_ids, position_ids=position_ids, head_mask=head_mask, inputs_embeds=inputs_embeds, output_attentions=output_attentions, output_hidden_states=output_hidden_states, return_dict=return_dict, ) sequence_output = outputs[0] logits = self.qa_outputs(sequence_output) start_logits, end_logits = logits.split(1, dim=-1) start_logits = start_logits.squeeze(-1).contiguous() end_logits = end_logits.squeeze(-1).contiguous() total_loss = None if start_positions is not None and end_positions is not None: # If we are on multi-GPU, split add a dimension if len(start_positions.size()) > 1: start_positions = start_positions.squeeze(-1) if len(end_positions.size()) > 1: end_positions = end_positions.squeeze(-1) # sometimes the start/end positions are outside our model inputs, we ignore these terms ignored_index = start_logits.size(1) start_positions = start_positions.clamp(0, ignored_index) end_positions = end_positions.clamp(0, ignored_index) loss_fct = CrossEntropyLoss(ignore_index=ignored_index) start_loss = loss_fct(start_logits, start_positions) end_loss = loss_fct(end_logits, end_positions) total_loss = (start_loss + end_loss) / 2 if not return_dict: output = (start_logits, end_logits) + outputs[2:] return ((total_loss,) + output) if total_loss is not None else output return QuestionAnsweringModelOutput( loss=total_loss, start_logits=start_logits, end_logits=end_logits, hidden_states=outputs.hidden_states, attentions=outputs.attentions, )
这个多头模型的“forward”函数会先调用“self.bert.forward()”,然后再调用“self.masked_lm_head.__call__()”方法来生成最终的结果。
- 预训练模型基类NewBertPreTrainedModel
- 完成了自定义模型类的编写后,可以使用该类创建一个模型实例:
newbert = NewBertForQuestionAnswering(newbert1_config)
模型权重可以通过调用“.from_pretrained()”加载:newbert.from_pretrained(pretrained_model_name_or_path="./您的权重文件本地存储路径/.")
- 编写自定义配置类。
后续操作
自定义模型文件构建完成后,可以参考托管模型到AI Gallery将模型文件托管至AI Gallery。建议托管的模型文件列表参见表2。
文件名称 |
描述 |
---|---|
config.json |
模型配置文件。 |
model.safetensors或pytorch_model.bin |
预训练模型的权重文件。 |
tokenizer.json |
(可选)预处理器的词表文件,用于初始化Tokenizer。 |
tokenizer_config.json |
(可选)预处理器的配置文件。 |
modeling_xxx.py |
(可选)自定义模型的代码文件,继承自PretrainedModel,包含实现自定义推理逻辑的代码。 |
configuration_xxx.py |
(可选)自定义配置的代码文件,继承自PretrainedConfig,包含实现自定义配置的逻辑代码。 |