如何在Amazon SageMaker上将TensorFlow checkpoint保存到S3？解决OSError报错

针对你遇到的OSError（无法写入S3路径的hdf5文件），以下是几个可行的解决办法：

1. 安装依赖让h5py支持S3路径写入 #

TensorFlow 2.6搭配的h5py本身不直接支持S3文件系统，需要通过fsspec和s3fs实现适配：

• 在训练环境中安装依赖：

  pip install fsspec s3fs
  

• 保持ModelCheckpoint的S3路径写法不变，此时h5py会通过fsspec自动处理S3的读写：

  from tensorflow.keras.callbacks import ModelCheckpoint
  
  checkpoint_callback = ModelCheckpoint(
      's3://mybucket/01-0.70.hdf5',
      save_best_only=True
  )
  model.fit(callbacks=[checkpoint_callback], ...)
  

2. 先保存到本地再同步到S3 #

绕开直接写入S3的兼容性问题，先将检查点保存到SageMaker的本地临时目录，再同步到S3：

from tensorflow.keras.callbacks import ModelCheckpoint
import boto3
import os

# 本地检查点路径
local_checkpoint_dir = '/tmp/checkpoints'
os.makedirs(local_checkpoint_dir, exist_ok=True)
checkpoint_callback = ModelCheckpoint(
    os.path.join(local_checkpoint_dir, '{epoch}-{val_loss:.2f}.hdf5'),
    save_best_only=True
)

# 训练模型
model.fit(callbacks=[checkpoint_callback], ...)

# 同步到S3
s3_client = boto3.client('s3')
for root, _, files in os.walk(local_checkpoint_dir):
    for file in files:
        local_path = os.path.join(root, file)
        s3_key = os.path.basename(local_path)
        s3_client.upload_file(local_path, 'mybucket', s3_key)

注意：/tmp目录有空间限制，若模型较大，可改用SageMaker挂载的EBS路径（如/opt/ml/checkpoints/）。

3. 改用TensorFlow原生SavedModel格式 #

TensorFlow对SavedModel格式的S3读写支持更完善，无需额外依赖，只需修改save_format参数：

checkpoint_callback = ModelCheckpoint(
    's3://mybucket/checkpoints/',
    save_format='tf',  # 使用SavedModel格式
    save_best_only=True,
    verbose=1
)
model.fit(callbacks=[checkpoint_callback], ...)

这种方式会在S3路径下生成SavedModel格式的目录结构，而非单一hdf5文件。

4. 确认IAM角色权限 #

确保SageMaker训练作业使用的IAM角色具备目标S3桶的读写权限，添加以下策略到角色中：

{
    "Version": "2012-10-17",
    "Statement": [
        {
            "Effect": "Allow",
            "Action": [
                "s3:PutObject",
                "s3:GetObject",
                "s3:ListBucket"
            ],
            "Resource": [
                "arn:aws:s3:::mybucket",
                "arn:aws:s3:::mybucket/*"
            ]
        }
    ]
}

内容的提问来源于stack exchange，提问作者Wt.N