Saving Pytorch model.state_dict() to s3

Question

I am trying to save a trained Pytorch model to S3. However, the torch.save(model.state_dict(), file_name) seems to support only local files. How can the state dict be saved to an S3 file?

I'm using Torch 0.4.0

Answer 1

As discussed by Soumith Chintala, Pytorch doesn't have custom APIs to do this job. However you can use boto3 or Petastorm library to solve the problem.

Here's a concrete example to write to an S3 object directly:

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')