读写 Checkpoint

更新时间：2026-05-14

bostorchconnector 支持直接对存储在 BOS 上的 checkpoint 做读写。

以 resnet 在某次 epoch 的结果写入和加载 checkpoint 为例：

                Python
                
            

                from bostorchconnector import BosCheckpoint, BosClientConfig

import torchvision
import torch

# 填充 <BUCKET>、<KEY>和对应的endpoint
CHECKPOINT_URI="bos://<BUCKET>/<KEY>/"
ENDPOINT="http://bj.bcebos.com"
config = BosClientConfig(log_level=1)
checkpoint = BosCheckpoint(endpoint=ENDPOINT, bos_client_config=config)

model = torchvision.models.resnet18()

# 保存checkpoint到Bos
with checkpoint.writer(CHECKPOINT_URI + "epoch0.ckpt") as writer:
    torch.save(model.state_dict(), writer)

# 从Bos读取checkpoint
with checkpoint.reader(CHECKPOINT_URI + "epoch0.ckpt") as reader:
    state_dict = torch.load(reader)

model.load_state_dict(state_dict)
            

分布式Checkpoint(DCP)

bostorchconnector 提供了对 PyTorch 分布式 Checkpoint 的支持，包括：

BosStorageWriter：实现了 PyTorch 的 StorageWriter 接口。
BosStorageReader：实现了 PyTorch 的 StorageReader 接口。
BosFileSystem：实现了 PyTorch 的 FileSystemBase 接口。

这些工具实现了 Bos 与 PyTorch 分布式 Checkpoint 的无缝集成，支持高效存储和读取分布式模型 Checkpoint。

前置条件与安装

需要 PyTorch 2.3 或更新版本。安装时需要指定 dcp 额外依赖：

                Shell
                
                pip install bostorchconnector[dcp]

示例

                Python
                
            

                from bostorchconnector.dcp import BosStorageWriter, BosStorageReader
from bostorchconnector import BosClientConfig 

import torchvision
import torch.distributed.checkpoint as DCP

# 配置
CHECKPOINT_URI = "bos://<BUCKET>/<KEY>/"
ENDPOINT = "http://bj.bcebos.com"

model = torchvision.models.resnet18()

# 自定义配置参数：credentials_path、log_level、log_path、part_size、pool_threads_num、max_attempts
config = BosClientConfig(part_size = 16 * 1024 * 1024, pool_threads_num = 64)

# 保存分布式 Checkpoint 到 Bos
bos_storage_writer = BosStorageWriter(
    endpoint=ENDPOINT,
    path=CHECKPOINT_URI,
    bos_client_config=config, # 可选
    thread_count=4,  # 可选，写入时使用的 IO 线程数
    overwrite=True,
)
DCP.save(
    state_dict=model.state_dict(),
    storage_writer=bos_storage_writer,
)

# 从 Bos 加载分布式 Checkpoint
model = torchvision.models.resnet18()
model_state_dict = model.state_dict()
bos_storage_reader = BosStorageReader(
    endpoint=ENDPOINT,
    path=CHECKPOINT_URI,
    bos_client_config=config, # 可选
)
DCP.load(
    state_dict=model_state_dict,
    storage_reader=bos_storage_reader,
)
model.load_state_dict(model_state_dict)
            

PyTorch Lightning

bostorchconnector 包含了对 PyTorch Lightning 的集成，提供了 BosLightningCheckpoint，它实现了 Lightning 的 CheckpointIO 接口。用户可以借此在 PyTorch Lightning 中使用 Bos 进行 Checkpoint 的读写。

安装

                Shell
                
                pip install bostorchconnector[lightning]

示例

                Python
                
            

                from lightning import Trainer
from bostorchconnector.lightning import BosLightningCheckpoint

from fsspec.registry import register_implementation
import bosfs

# ...

CHECKPOINT_URI = "bos://<BUCKET>/<KEY>/"
ENDPOINT = "http://bj.bcebos.com"

save_only_latest = True

register_implementation("bos", bosfs.BOSFileSystem)

dataset = WikiText2()
dataloader = DataLoader(dataset, num_workers=2)

model = LightningTransformer(vocab_size=dataset.vocab_size)
bos_lightning_checkpoint = BosLightningCheckpoint(endpoint=ENDPOINT)

checkpoint_callback = ModelCheckpoint(
    dirpath=CHECKPOINT_URI,
    save_top_k=1 if save_only_latest else -1,
    every_n_train_steps=1,
    filename="checkpoint-{epoch:02d}-{step:02d}",
    enable_version_counter=True,
)

trainer = Trainer(
    plugins=[bos_lightning_checkpoint],
    callbacks=[checkpoint_callback],
    min_epochs=4,
    max_epochs=5,
    max_steps=3,
)
trainer.fit(model, dataloader)

# read
r_trainer = Trainer(
    plugins=[bos_lightning_checkpoint],
    min_epochs=4,
    max_epochs=5,
    max_steps=3,
)
# Load the checkpoint in `ckpt_path` before training
r_trainer.fit(model, dataloader, ckpt_path=CHECKPOINT_URI + "checkpoint-epoch=00-step=03.ckpt")
            

评价此篇文章

有帮助没帮助

构建 Dataset

性能测试

百度智能云

BOS 对象存储

BOS 对象存储

读写 Checkpoint

分布式Checkpoint(DCP)

前置条件与安装

示例

PyTorch Lightning

安装

示例