35 lines
926 B
Python
35 lines
926 B
Python
"""Test that a model can be saved and loaded again."""
|
|
|
|
from pathlib import Path
|
|
|
|
import torch
|
|
from dotenv import load_dotenv
|
|
from torchinfo import summary
|
|
|
|
from model.src.models import VisualCommunicationModel
|
|
from shared.data_store import connect_minio, put_model
|
|
from shared.utils import setup_logging
|
|
|
|
DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
|
|
|
|
if __name__ == '__main__':
|
|
# load in env file
|
|
env_path = Path(__file__).parent.parent / 'server.env'
|
|
assert env_path.exists()
|
|
load_dotenv(env_path)
|
|
# setup logging
|
|
setup_logging()
|
|
# connect to minio
|
|
client = connect_minio()
|
|
# instantiate model
|
|
model = VisualCommunicationModel().to(DEVICE)
|
|
# show model weights
|
|
summary(model)
|
|
# put buffer in minio bucket
|
|
hash_str = put_model(
|
|
client=client,
|
|
model=model,
|
|
)
|
|
print(f"hash string: {hash_str}")
|
|
print('saved to Minio')
|