基于PyTorch-Lightning的LayoutLMv3训练：同步Checkpoint到HF及配置

1. 每个Epoch结束自动上传Checkpoint到Hugging Face #

自定义PyTorch-Lightning回调，监听验证阶段结束事件，自动将最新保存的Checkpoint上传至Hugging Face仓库。

步骤1：实现上传回调 #

from pytorch_lightning.callbacks import Callback
from huggingface_hub import upload_file, login

# 执行一次登录Hugging Face（后续无需重复执行）
login()

class HFCheckpointUploader(Callback):
    def __init__(self, repo_name):
        self.repo_name = repo_name

    def on_validation_end(self, trainer, pl_module):
        # 获取最新生成的Checkpoint路径
        latest_ckpt = trainer.checkpoint_callback.last_model_path
        if latest_ckpt:
            # 上传至仓库的checkpoints目录下，保留原文件名
            upload_file(
                path_or_fileobj=latest_ckpt,
                path_in_repo=f"checkpoints/{latest_ckpt.split('/')[-1]}",
                repo_id=self.repo_name
            )
            print(f"Checkpoint已上传至: {self.repo_name}/checkpoints/")

步骤2：添加到训练回调列表 #

# 初始化上传回调，替换为你的Hugging Face仓库名
hf_uploader = HFCheckpointUploader(repo_name="your-repo-name")

# 更新Trainer配置，加入上传回调
trainer = pl.Trainer(
    accelerator="gpu",
    precision=16,
    max_epochs=50,
    callbacks=[
        model_checkpoint,
        early_stopping,
        hf_uploader  # 新增上传回调
    ],
)

2. 从Hugging Face加载最佳模型 #

提供两种加载方式，适配不同使用场景：

方式一：加载完整Lightning Module Checkpoint #

先从Hub下载目标Checkpoint文件，再用原有逻辑加载为ModelModule：

from huggingface_hub import hf_hub_download

# 替换为你仓库中最佳Checkpoint的文件名
target_ckpt = "checkpoints/epoch-5-step-1234-val_loss-0.1234.ckpt"
checkpoint_path = hf_hub_download(
    repo_id="your-repo-name",
    filename=target_ckpt
)

# 加载为ModelModule
trained_model = ModelModule.load_from_checkpoint(
    checkpoint_path,
    n_classes=len(DOCUMENT_CLASSES)
)

方式二：直接加载预训练模型权重 #

如果之前已将trained_model.model推送到Hub，可直接加载核心模型组件：

from transformers import LayoutLMv3ForSequenceClassification, LayoutLMv3Processor

# 加载模型和处理器
model = LayoutLMv3ForSequenceClassification.from_pretrained("your-repo-name")
processor = LayoutLMv3Processor.from_pretrained("your-repo-name")

# 可选：包装为Lightning Module
trained_model = ModelModule(len(DOCUMENT_CLASSES))
trained_model.model = model

3. 在Hugging Face配置TensorBoard查看指标 #

方法一：训练时自动同步日志 #

使用HuggingFaceLogger，训练过程中自动上传TensorBoard日志到Hub：

from pytorch_lightning.loggers import HuggingFaceLogger

hf_logger = HuggingFaceLogger(
    repo_name="your-repo-name",
    log_model="all"  # 可选：自动记录模型版本
)

# 更新Trainer配置，加入日志器
trainer = pl.Trainer(
    accelerator="gpu",
    precision=16,
    max_epochs=50,
    callbacks=[model_checkpoint, early_stopping, hf_uploader],
    logger=hf_logger
)

训练完成后，直接在Hugging Face仓库页面的「TensorBoard」标签页查看指标图表。

方法二：手动上传本地日志 #

若已有本地lightning_logs/日志目录，手动上传至Hub：

from huggingface_hub import upload_folder

upload_folder(
    folder_path="lightning_logs/",
    path_in_repo="tensorboard_logs/",
    repo_id="your-repo-name"
)

在仓库页面点击「Add file」下拉菜单，选择「View TensorBoard」，指定tensorboard_logs/作为日志目录即可查看。

内容的提问来源于stack exchange，提问作者im_s_kumar