自定义模型规范

自定义模型的使用流程

1. 托管模型到AI Gallery。
   • 模型基础设置里的“任务类型”选择除“文本问答”和“文本生成”之外的类型。
   • 上传模型文件时需要确认待上传的文件是否满足自定义模型规范。如果模型要支持训练，则需要满足自定义模型规范（训练）；如果模型要支持推理，则需要满足自定义模型规范（推理）。

2. 发布模型到AI Gallery。
3. 使用AI Gallery微调大师训练模型或使用AI Gallery在线推理服务部署模型。
   • 如果进行模型微调，则“训练任务类型”选择“自定义”。
   • 如果部署为推理服务，则“推理任务类型”选择“自定义”

自定义模型规范（训练）

当托管自定义模型到AI Gallery时，如果模型要支持AI Gallery的模型微调，则需要在“模型文件”添加gallery_train文件夹，文件夹内容参考表1。

• gallery_train文件夹必须是一级目录直接上传，否则会被判定不符合自定义模型规范，无法使用模型微调。
• 如果自定义模型的模型文件不符合gallery_train文件列表要求或文件内容为空，都将不能正常进行模型微调。

自定义模型规范（推理）

当托管自定义模型到AI Gallery时，如果模型要支持AI Gallery的推理服务，则需要在“模型文件”添加gallery_inference文件夹，文件夹内容参考表2。

• gallery_inference文件夹必须是一级目录直接上传，否则会被判定不符合自定义模型规范，无法使用模型微调。
• 如果自定义模型的模型文件不符合gallery_inference文件列表要求或文件内容为空，都将不能正常部署在线推理服务。

自定义模型使用的预置镜像

AI Gallery提供了PyTorch基础镜像，镜像里已经安装好了运行任务所需的软件，供自定义模型直接使用，快速进行训练、推理。预置镜像的版本信息请参见表3。

“train.py”示例

import json
import os

from datasets import load_dataset
from transformers import AutoImageProcessor
from torchvision.transforms import RandomResizedCrop, Compose, Normalize, ToTensor, RandomHorizontalFlip
import numpy as np
from transformers import AutoModelForImageClassification, TrainingArguments, Trainer
from transformers import DefaultDataCollator
from sklearn import metrics

# 环境变量
# 工作目录
ENV_AG_WORK_DIR = 'ENV_AG_WORK_DIR'
# 模型存放路径
ENV_AG_MODEL_DIR = 'ENV_AG_MODEL_DIR'
# 数据集存放路径
ENV_AG_DATASET_DIR = 'ENV_AG_DATASET_DIR'
# 配置的训练超参json字符串
ENV_AG_USER_PARAMS = 'ENV_AG_USER_PARAMS'
# 训练产物存放路径
ENV_AG_TRAIN_OUTPUT_DIR = 'ENV_AG_TRAIN_OUTPUT_DIR'

_transforms = None


def _multi_class_classification_metrics(pred):
    raw_predictions, labels = pred
    predictions = np.argmax(raw_predictions, axis=1)
    results = {
        "f1_macro": metrics.f1_score(labels, predictions, average="macro"),
        "f1_micro": metrics.f1_score(labels, predictions, average="micro"),
        "f1_weighted": metrics.f1_score(labels, predictions, average="weighted"),
        "precision_macro": metrics.precision_score(labels, predictions, average="macro"),
        "precision_micro": metrics.precision_score(labels, predictions, average="micro"),
        "precision_weighted": metrics.precision_score(labels, predictions, average="weighted"),
        "recall_macro": metrics.recall_score(labels, predictions, average="macro"),
        "recall_micro": metrics.recall_score(labels, predictions, average="micro"),
        "recall_weighted": metrics.recall_score(labels, predictions, average="weighted"),
        "accuracy": metrics.accuracy_score(labels, predictions),
    }
    return results


def parse_args():
    """ 从AIGallery环境变量中获取用户配置的超参json """
    return json.loads(os.getenv(ENV_AG_USER_PARAMS))


def _process_input_data(image_processor):
    # 加载数据集
    dataset_path = os.getenv(ENV_AG_DATASET_DIR)
    dataset = load_dataset("imagefolder", data_dir=dataset_path)

    # 数据增强
    normalize = Normalize(mean=image_processor.image_mean, std=image_processor.image_std)
    size = (image_processor.size["shortest_edge"] if "shortest_edge" in image_processor.size else (
        image_processor.size["height"], image_processor.size["width"]))
    global _transforms
    _transforms = Compose([RandomResizedCrop(size), RandomHorizontalFlip(), ToTensor(), normalize])
    ret = dataset.with_transform(_format_transforms)
    return ret


# 转换函数
def _format_transforms(examples):
    examples["pixel_values"] = [_transforms(img.convert("RGB")) for img in examples["image"]]
    del examples["image"]
    return examples


def train(user_args):
    print('Start to process dataset')
    model_path = os.getenv(ENV_AG_MODEL_DIR)
    image_processor = AutoImageProcessor.from_pretrained(model_path)

    dataset = _process_input_data(image_processor)
    print(f"Dataset: {dataset}")
    # label和id映射
    classes = dataset["train"].features["label"].names
    label2id = {c: i for i, c in enumerate(classes)}
    id2label = {i: c for i, c in enumerate(classes)}

    print('Start to load model')
    # 加载模型
    model = AutoModelForImageClassification.from_pretrained(
        model_path,
        num_labels=len(classes),
        id2label=id2label,
        label2id=label2id,
        ignore_mismatched_sizes=True
    )

    print('Start to set training args')
    # 训练参数
    training_args = TrainingArguments(
        output_dir=os.getenv(ENV_AG_TRAIN_OUTPUT_DIR),
        remove_unused_columns=False,
        evaluation_strategy="epoch",
        save_strategy=user_args['save_strategy'],
        learning_rate=float(user_args['lr']),
        save_total_limit=3,
        per_device_train_batch_size=32,
        gradient_accumulation_steps=1,
        per_device_eval_batch_size=int(user_args['per_device_eval_batch_size']),
        num_train_epochs=int(user_args['num_train_epochs']),
        warmup_ratio=float(user_args['warmup_ratio']),
        logging_steps=int(user_args['logging_steps']),
        load_best_model_at_end=True,
        metric_for_best_model="accuracy",
        push_to_hub=False,
    )

    print('Start to train')
    # 训练参数
    trainer = Trainer(
        model=model,
        args=training_args,
        data_collator=DefaultDataCollator(),
        train_dataset=dataset["train"],
        eval_dataset=dataset["test"],
        tokenizer=image_processor,
        compute_metrics=_multi_class_classification_metrics,
    )

    # 开始训练
    train_results = trainer.train()
    print('Start to save model')
    # 保存模型
    trainer.save_model()
    trainer.log_metrics("train", train_results.metrics)
    trainer.save_metrics("train", train_results.metrics)
    trainer.save_state()

    print('Start to evaluate')
    # 在验证集上做准确性评估
    eva_metrics = trainer.evaluate()
    trainer.log_metrics("eval", eva_metrics)
    trainer.save_metrics("eval", eva_metrics)

    print('All Done')


if __name__ == '__main__':
    args = parse_args()
    train(args)

“train_params.json”示例

{
    "training_methods": [
        {
            "name": "全参微调",
            "hyperparameters": [
                {
                    "name": "lr",
                    "type": "float",
                    "required": true,
                    "default": 0.001,
                    "help": "学习率"
                },
                {
                    "name": "per_device_eval_batch_size",
                    "type": "int",
                    "required": false,
                    "default": 32,
                    "help": "批大小"
                },
                {
                    "name": "logging_steps",
                    "type": "int",
                    "required": false,
                    "default": 24,
                    "help": "每多少步记录一次步骤"
                },
                {
                    "name": "save_strategy",
                    "type": "str",
                    "required": true,
                    "default": "epoch",
                    "help": "训练过程中保存checkpoint的策略"
                },
                {
                    "name": "num_train_epochs",
                    "type": "int",
                    "required": true,
                    "default": 20,
                    "help": "训练的总epochs数"
                },
                {
                    "name": "warmup_ratio",
                    "type": "float",
                    "required": true,
                    "default": 0.1,
                    "help": "用于指定线性热身占总训练步骤的比例"
                }
            ]
        }
    ]
}

“inference.py”示例

from typing import Dict, List, Any
from transformers import pipeline
import os


class EndpointHandler:
    def __init__(self, path=""):
        # Use a pipeline as a high-level helper
        self.pipe = pipeline("question-answering", model=path)

    def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
        """
         data args:
              inputs (:obj: `str`)
        Return:
              A :obj:`list` | `dict`: will be serialized and returned
        """
        # get inputs
        inputs = data.pop("inputs", data)
        question = inputs["question"]
        context = inputs["context"]
        resp = self.pipe(question=question, context=context)
        return resp