tests/model_io_test.py

@@ -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')