“train.py”示例

import jsonimport osfrom datasets import load_datasetfrom transformers import AutoImageProcessorfrom torchvision.transforms import RandomResizedCrop, Compose, Normalize, ToTensor, RandomHorizontalFlipimport numpy as npfrom transformers import AutoModelForImageClassification, TrainingArguments, Trainerfrom transformers import DefaultDataCollatorfrom 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 = Nonedef _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 resultsdef 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 examplesdef 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)