python - 将 Pytorch model.state_dict() 保存到 s3

Question

我正在尝试将经过训练的 Pytorch 模型保存到 S3。然而，torch.save(model.state_dict(), file_name)似乎只支持本地文件。如何将状态字典保存到 S3 文件中？

我正在使用 Torch 0.4.0

Answer 1

正如 Soumith Chintala 所讨论的, Pytorch 没有自定义 API 来完成这项工作。但是你可以使用 boto3 或 Petastorm 库来解决这个问题。

这是一个直接写入 S3 对象的具体示例:

import boto3

# Convert your existing model to JSON
saved_model = model.to_json()

# Write JSON object to S3 as "model.json"
client = boto3.client('s3')
client.put_object(Body=saved_model,
                  Bucket='BUCKET_NAME',
                  Key='model.json')