46 lines
1.2 KiB
Python
46 lines
1.2 KiB
Python
"""Main script to be run by service."""
|
|
from __future__ import annotations
|
|
|
|
from hashlib import md5
|
|
from io import BytesIO
|
|
|
|
import torch
|
|
from data_store import connect
|
|
from data_store import put
|
|
from models import VisualCommunicationModel
|
|
from torchinfo import summary
|
|
from utils import check_env
|
|
|
|
from shared.utils import setup_logging
|
|
|
|
|
|
DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
|
|
|
|
if __name__ == '__main__':
|
|
# check that environment variables are set
|
|
check_env()
|
|
# setup logging
|
|
setup_logging()
|
|
# instantiate model
|
|
model = VisualCommunicationModel().to(DEVICE)
|
|
# show model weights
|
|
summary(model)
|
|
# for param in model.parameters():
|
|
# logging.info(param.data)
|
|
# save model to buffer
|
|
buffer = BytesIO()
|
|
torch.save(model.state_dict(), buffer)
|
|
print(f"buffer size: {len(buffer.getvalue())}")
|
|
# calculate buffer hash
|
|
hash_str = md5(buffer.getbuffer()).hexdigest()
|
|
print(f"hash string: {hash_str}")
|
|
# connect to minio
|
|
client = connect()
|
|
# put buffer in minio bucket
|
|
put(
|
|
client=client,
|
|
buffer=buffer,
|
|
object_name=hash_str,
|
|
)
|
|
print('saved to Minio')
|