added model instantiation test
This commit is contained in:
@@ -0,0 +1,37 @@
|
||||
"""Test that a model can be saved and loaded again."""
|
||||
|
||||
from io import BytesIO
|
||||
|
||||
import torch
|
||||
from torchinfo import summary
|
||||
|
||||
from model.src.models import VisualCommunicationModel
|
||||
from shared.data_store import connect, put
|
||||
from shared.utils import check_env, 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())}")
|
||||
# connect to minio
|
||||
client = connect(bucket='models')
|
||||
# put buffer in minio bucket
|
||||
hash_str = put(
|
||||
client=client,
|
||||
buffer=buffer,
|
||||
)
|
||||
print(f"hash string: {hash_str}")
|
||||
print('saved to Minio')
|
||||
Reference in New Issue
Block a user