123 lines
6.7 KiB
Python
123 lines
6.7 KiB
Python
import os
|
||
import datasets
|
||
|
||
from pathlib import Path
|
||
from transformers import (
|
||
ResNetConfig,
|
||
TrainingArguments,
|
||
Trainer
|
||
)
|
||
|
||
from ..utils import preprocess_fn
|
||
from ..model.Resizer import Resizer
|
||
from ...globals import NUM_CHANNELS, NUM_CLASSES, RESIZER_IMG_SIZE
|
||
|
||
|
||
def train():
|
||
cur_dirpath = os.getcwd()
|
||
script_dirpath = Path(__file__).resolve().parent
|
||
os.chdir(script_dirpath)
|
||
|
||
data = datasets.load_dataset("./dataset").shuffle(seed=42)
|
||
data = data.rename_column("images", "pixel_values")
|
||
data.flatten_indices()
|
||
data = data.with_transform(preprocess_fn)
|
||
train_data, test_data = data['train'], data['test']
|
||
|
||
config = ResNetConfig(
|
||
num_channels=NUM_CHANNELS,
|
||
num_labels=NUM_CLASSES,
|
||
img_size=RESIZER_IMG_SIZE
|
||
)
|
||
model = Resizer(config)
|
||
model = Resizer.from_pretrained("/home/lhy/code/TeXify/src/models/resizer/train/train_result_pred_height_v4/checkpoint-213000")
|
||
|
||
training_args = TrainingArguments(
|
||
# resume_from_checkpoint="/home/lhy/code/TeXify/src/models/resizer/train/train_result_pred_height_v3/checkpoint-94500",
|
||
max_grad_norm=1.0,
|
||
# use_cpu=True,
|
||
seed=42, # 随机种子，用于确保实验的可重复性
|
||
# data_seed=42, # data sampler的采样也固定
|
||
# full_determinism=True, # 使整个训练完全固定（这个设置会有害于模型训练，只用于debug）
|
||
|
||
output_dir='./train_result_pred_height_v5', # 输出目录
|
||
overwrite_output_dir=False, # 如果输出目录存在，不删除原先的内容
|
||
report_to=["tensorboard"], # 输出日志到TensorBoard，
|
||
#+通过在命令行：tensorboard --logdir ./logs 来查看日志
|
||
|
||
logging_dir=None, # TensorBoard日志文件的存储目录
|
||
log_level="info",
|
||
logging_strategy="steps", # 每隔一定步数记录一次日志
|
||
logging_steps=500, # 记录日志的步数间隔
|
||
logging_nan_inf_filter=False, # 对loss=nan或inf进行记录
|
||
|
||
num_train_epochs=50, # 总的训练轮数
|
||
# max_steps=3, # 训练的最大步骤数。如果设置了这个参数，
|
||
#+那么num_train_epochs将被忽略（通常用于调试）
|
||
|
||
# label_names = ['your_label_name'], # 指定data_loader中的标签名，如果不指定则默认为'labels'
|
||
|
||
per_device_train_batch_size=55, # 每个GPU的batch size
|
||
per_device_eval_batch_size=48*2, # 每个GPU的evaluation batch size
|
||
auto_find_batch_size=False, # 自动搜索合适的batch size（指数decay）
|
||
|
||
optim = 'adamw_torch', # 还提供了很多AdamW的变体（相较于经典的AdamW更加高效）
|
||
#+当设置了optim后，就不需要在Trainer中传入optimizer
|
||
lr_scheduler_type="cosine", # 设置lr_scheduler
|
||
warmup_ratio=0.1, # warmup占整个训练steps的比例
|
||
# warmup_steps=500, # 预热步数
|
||
weight_decay=0, # 权重衰减
|
||
learning_rate=5e-5, # 学习率
|
||
fp16=False, # 是否使用16位浮点数进行训练
|
||
gradient_accumulation_steps=1, # 梯度累积步数，当batch size无法开很大时，可以考虑这个参数来实现大batch size的效果
|
||
gradient_checkpointing=False, # 当为True时，会在forward时适当丢弃一些中间量（用于backward），从而减轻显存压力（但会增加forward的时间）
|
||
label_smoothing_factor=0.0, # softlabel，等于0时表示未开启
|
||
# debug='underflow_overflow', # 训练时检查溢出，如果发生，则会发出警告。（该模式通常用于debug）
|
||
torch_compile=True, # 是否使用torch.compile来编译模型（从而获得更好的训练和推理性能）
|
||
#+ 要求torch > 2.0，并且这个功能现在还不是很稳定
|
||
# deepspeed='your_json_path', # 使用deepspeed来训练，需要指定ds_config.json的路径
|
||
#+ 在Trainer中使用Deepspeed时一定要注意ds_config.json中的配置是否与Trainer的一致（如学习率，batch size，梯度累积步数等）
|
||
#+ 如果不一致，会出现很奇怪的bug（而且一般还很难发现）
|
||
|
||
dataloader_pin_memory=True, # 可以加快数据在cpu和gpu之间转移的速度
|
||
dataloader_num_workers=16, # 默认不会使用多进程来加载数据
|
||
dataloader_drop_last=True, # 丢掉最后一个minibatch
|
||
|
||
evaluation_strategy="steps", # 评估策略，可以是"steps"或"epoch"
|
||
eval_steps=500, # if evaluation_strategy="step"
|
||
# eval_steps=10, # if evaluation_strategy="step"
|
||
|
||
save_strategy="steps", # 保存checkpoint的策略
|
||
save_steps=1500, # 模型保存的步数间隔
|
||
save_total_limit=5, # 保存的模型的最大数量。如果超过这个数量，最旧的模型将被删除
|
||
|
||
load_best_model_at_end=True, # 训练结束时是否加载最佳模型
|
||
metric_for_best_model="eval_loss", # 用于选择最佳模型的指标
|
||
greater_is_better=False, # 指标值越小越好
|
||
|
||
do_train=True, # 是否进行训练，通常用于调试
|
||
do_eval=True, # 是否进行评估，通常用于调试
|
||
|
||
remove_unused_columns=True, # 是否删除没有用到的列（特征），默认为True
|
||
#+当删除了没用到的列后，making it easier to unpack inputs into the model’s call function
|
||
|
||
push_to_hub=False, # 是否训练完后上传hub，需要先在命令行：huggingface-cli login进行登录认证的配置，配置完后，认证信息会存到cache文件夹里
|
||
hub_model_id="a_different_name", # 模型的名字
|
||
#+每次保存模型时，都会上传到hub，
|
||
#+训练完后，记得trainer.push_to_hub()，会将模型使用的参数以及验证集上的结果传到hub上
|
||
)
|
||
|
||
trainer = Trainer(
|
||
model,
|
||
training_args,
|
||
train_dataset=train_data,
|
||
eval_dataset=test_data,
|
||
)
|
||
trainer.train()
|
||
|
||
os.chdir(cur_dirpath)
|
||
|
||
|
||
if __name__ == '__main__':
|
||
train()
|