“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)