可以优化的地方:
per_device_train_batch_size(相当于batch size,越小显存占的越小)
gradient_accumulation_steps(per_device_train_batch_size*gradient_accumulation_steps=计算梯度的数据数)
gradient_checkpointing(前项激活值里面有很多是不需要存的,可以在反向传播再次计算的)
optim(可以改为adafactor)
冻结参数(只训练下游任务的参数)
将max_length减小
参考代码:
train_args = TrainingArguments(output_dir="./checkpoints", # 输出文件夹
per_device_train_batch_size=1, # 训练时的batch_size
gradient_accumulation_steps=32, # *** 梯度累加 ***
gradient_checkpointing=True, # *** 梯度检查点 *** 前项激活值里面有很多是不需要存的,可以在反向传播再次计算的
optim="adafactor", # *** adafactor优化器 ***
per_device_eval_batch_size=1, # 验证时的batch_size
num_train_epochs=1, # 训练轮数
logging_steps=10, # log 打印的频率
evaluation_strategy="epoch", # 评估策略
save_strategy="epoch", # 保存策略
save_total_limit=3, # 最大保存数
learning_rate=2e-5, # 学习率
weight_decay=0.01, # weight_decay
metric_for_best_model="f1", # 设定评估指标
load_best_model_at_end=True) # 训练完成后加载最优模型
for name, param in model.bert.named_parameters():
param.requires_grad = False
tokenized_examples = tokenizer(examples["review"], max_length=32, truncation=True, padding="max_length")