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