import argparse import functools import os from peft import LoraConfig, get_peft_model, AdaLoraConfig, PeftModel, prepare_model_for_kbit_training from transformers import Seq2SeqTrainer, Seq2SeqTrainingArguments, WhisperForConditionalGeneration, WhisperProcessor from utils.callback import SavePeftModelCallback from utils.data_utils import DataCollatorSpeechSeq2SeqWithPadding from utils.model_utils import load_from_checkpoint from utils.reader import CustomDataset from utils.utils import print_arguments, make_inputs_require_grad, add_arguments parser = argparse.ArgumentParser(description=__doc__) add_arg = functools.partial(add_arguments, argparser=parser) add_arg("train_data", type=str, default="dataset/train.json", help="训练数据集的路径") add_arg("test_data", type=str, default="dataset/test.json", help="测试数据集的路径") add_arg("base_model", type=str, default="openai/whisper-tiny", help="Whisper的基础模型") add_arg("output_dir", type=str, default="output/", help="训练保存模型的路径") add_arg("warmup_steps", type=int, default=50, help="训练预热步数") add_arg("logging_steps", type=int, default=100, help="打印日志步数") add_arg("eval_steps", type=int, default=1000, help="多少步数评估一次") add_arg("save_steps", type=int, default=1000, help="多少步数保存模型一次") add_arg("num_workers", type=int, default=8, help="读取数据的线程数量") add_arg("learning_rate", type=float, default=1e-3, help="学习率大小") add_arg("min_audio_len", type=float, default=0.5, help="最小的音频长度，单位秒") add_arg("max_audio_len", type=float, default=30, help="最大的音频长度，单位秒，不能大于30秒") add_arg("use_adalora", type=bool, default=True, help="是否使用AdaLora而不是Lora") add_arg("fp16", type=bool, default=True, help="是否使用fp16训练模型") add_arg("use_8bit", type=bool, default=False, help="是否将模型量化为8位") add_arg("timestamps", type=bool, default=False, help="训练时是否使用时间戳数据") add_arg("use_compile", type=bool, default=False, help="是否使用Pytorch2.0的编译器") add_arg("local_files_only", type=bool, default=False, help="是否只在本地加载模型，不尝试下载") add_arg("num_train_epochs", type=int, default=3, help="训练的轮数") add_arg("language", type=str, default="Chinese", help="设置语言，可全称也可简写，如果为None则训练的是多语言") add_arg("task", type=str, default="transcribe", choices=['transcribe', 'translate'], help="模型的任务") add_arg("augment_config_path", type=str, default=None, help="数据增强配置文件路径") add_arg("resume_from_checkpoint", type=str, default=None, help="恢复训练的检查点路径") add_arg("per_device_train_batch_size", type=int, default=8, help="训练的batch size") add_arg("per_device_eval_batch_size", type=int, default=8, help="评估的batch size") add_arg("gradient_accumulation_steps", type=int, default=1, help="梯度累积步数") add_arg("push_to_hub", type=bool, default=False, help="是否将模型权重推到HuggingFace Hub") add_arg("hub_model_id", type=str, default=None, help="HuggingFace Hub上的模型仓库ID") args = parser.parse_args() print_arguments(args) def main(): # 获取Whisper的数据处理器，这个包含了特征提取器、tokenizer processor = WhisperProcessor.from_pretrained(args.base_model, language=args.language, task=args.task, no_timestamps=not args.timestamps, local_files_only=args.local_files_only) # 读取数据 train_dataset = CustomDataset(data_list_path=args.train_data, processor=processor, language=args.language, timestamps=args.timestamps, min_duration=args.min_audio_len, max_duration=args.max_audio_len, augment_config_path=args.augment_config_path) test_dataset = CustomDataset(data_list_path=args.test_data, processor=processor, language=args.language, timestamps=args.timestamps, min_duration=args.min_audio_len, max_duration=args.max_audio_len) print(f"训练数据：{len(train_dataset)}，测试数据：{len(test_dataset)}") # 数据padding器 data_collator = DataCollatorSpeechSeq2SeqWithPadding(processor=processor) # 获取Whisper模型 device_map = "auto" world_size = int(os.environ.get("WORLD_SIZE", 1)) ddp = world_size != 1 if ddp: device_map = {"": int(os.environ.get("LOCAL_RANK") or 0)} # 获取模型 model = WhisperForConditionalGeneration.from_pretrained(args.base_model, load_in_8bit=args.use_8bit, device_map=device_map, local_files_only=args.local_files_only) model.config.forced_decoder_ids = None model.config.suppress_tokens = [] # 量化模型 model = prepare_model_for_kbit_training(model) # 注册forward，否则多卡训练会失败 model.model.encoder.conv1.register_forward_hook(make_inputs_require_grad) print('加载LoRA模块...') if args.resume_from_checkpoint: # 恢复训练时加载Lora参数 print("Loading adapters from checkpoint.") model = PeftModel.from_pretrained(model, args.resume_from_checkpoint, is_trainable=True) else: print(f'adding LoRA modules...') target_modules = ["k_proj", "q_proj", "v_proj", "out_proj", "fc1", "fc2"] print(target_modules) if args.use_adalora: config = AdaLoraConfig(init_r=12, target_r=4, beta1=0.85, beta2=0.85, tinit=200, tfinal=1000, deltaT=10, lora_alpha=32, lora_dropout=0.1, orth_reg_weight=0.5, target_modules=target_modules) else: config = LoraConfig(r=32, lora_alpha=64, target_modules=target_modules, lora_dropout=0.05, bias="none") model = get_peft_model(model, config) if args.base_model.endswith("/"): args.base_model = args.base_model[:-1] output_dir = os.path.join(args.output_dir, os.path.basename(args.base_model)) # 定义训练参数 training_args = \ Seq2SeqTrainingArguments(output_dir=output_dir, # 保存检查点和意志的目录 per_device_train_batch_size=args.per_device_train_batch_size, # 训练batch_size大小 per_device_eval_batch_size=args.per_device_eval_batch_size, # 评估batch_size大小 gradient_accumulation_steps=args.gradient_accumulation_steps, # 训练梯度累计步数 learning_rate=args.learning_rate, # 学习率大小 warmup_steps=args.warmup_steps, # 预热步数 num_train_epochs=args.num_train_epochs, # 微调训练轮数 save_strategy="steps", # 指定按照步数保存检查点 evaluation_strategy="steps", # 指定按照步数评估模型 load_best_model_at_end=True, # 指定是否在结束时加载最优模型 fp16=args.fp16, # 是否使用半精度训练 report_to=["tensorboard"], # 指定使用tensorboard保存log save_steps=args.save_steps, # 指定保存检查点的步数 eval_steps=args.eval_steps, # 指定评估模型的步数 torch_compile=args.use_compile, # 使用Pytorch2.0的编译器 save_total_limit=5, # 只保存最新检查点的数量 optim='adamw_torch', # 指定优化方法 ddp_find_unused_parameters=False if ddp else None, # 分布式训练设置 dataloader_num_workers=args.num_workers, # 设置读取数据的线程数量 logging_steps=args.logging_steps, # 指定打印log的步数 remove_unused_columns=False, # 删除模型不需要的数据列 label_names=["labels"], # 与标签对应的输入字典中的键列表 push_to_hub=args.push_to_hub, ) if training_args.local_rank == 0 or training_args.local_rank == -1: print('=' * 90) model.print_trainable_parameters() print('=' * 90) # 定义训练器 trainer = Seq2SeqTrainer(args=training_args, model=model, train_dataset=train_dataset, eval_dataset=test_dataset, data_collator=data_collator, tokenizer=processor.feature_extractor, callbacks=[SavePeftModelCallback]) model.config.use_cache = False trainer._load_from_checkpoint = load_from_checkpoint # 开始训练 trainer.train(resume_from_checkpoint=args.resume_from_checkpoint) # 保存最后的模型 trainer.save_state() # 重新启用缓存以更快地推断 model.config.use_cache = True if training_args.local_rank == 0 or training_args.local_rank == -1: model.save_pretrained(os.path.join(output_dir, "checkpoint-final")) # 是否把模型参数文件推送到huggingface if training_args.push_to_hub: hub_model_id = args.hub_model_id if args.hub_model_id is not None else output_dir model.push_to_hub(hub_model_id) if __name__ == '__main__': main()