import torch | |
import io | |
import json | |
from petrel_client.client import Client | |
client = Client('./conf/petreloss.conf') | |
# Pytorch save & load | |
data = torch.tensor([0, 1, 2, 3]) | |
tensor_url = 's3://bucket1/tensor_data' | |
with io.BytesIO() as f: | |
torch.save(data, f) | |
client.put(tensor_url, f.getvalue()) | |
with io.BytesIO(client.get(tensor_url)) as f: | |
data2 = torch.load(f) | |
assert torch.equal(data, data2) | |
# Json dumps & loads | |
data = [0, 1, 2, 3] | |
json_data_url = 's3://bucket1/json_data' | |
client.put(json_data_url, json.dumps(data).encode('utf-8')) | |
data2 = json.loads(client.get(json_data_url)) | |
assert data == data2 | |