examples

Author: Borda

examples/train-callback.py

@@ -4,14 +4,18 @@
 from litmodels import upload_model
 from sample_model import LitAutoEncoder
 
+# Define the model name - this should be unique to your model
+# The format is <organization>/<teamspace>/<model-name>
+MY_MODEL_NAME = "jirka/kaggle/lit-auto-encoder-callback"
+
 
 class UploadModelCallback(Callback):
     def on_train_epoch_end(self, trainer, pl_module):
         # Get the best model path from the checkpoint callback
         best_model_path = trainer.checkpoint_callback.best_model_path
         if best_model_path:
             print(f"Uploading model: {best_model_path}")
-            upload_model(path=best_model_path, name="jirka/kaggle/lit-auto-encoder-callback")
+            upload_model(path=best_model_path, name=MY_MODEL_NAME)
 
 
 if __name__ == "__main__":

examples/train-resume.py

@@ -4,11 +4,16 @@
 from litmodels import download_model
 from sample_model import LitAutoEncoder
 
+# Define the model name - this should be unique to your model
+# The format is <organization>/<teamspace>/<model-name>:<model-version>
+MY_MODEL_NAME = "jirka/kaggle/lit-auto-encoder-callback:latest"
+
+
 if __name__ == "__main__":
     dataset = tv.datasets.MNIST(".", download=True, transform=tv.transforms.ToTensor())
     train, val = data.random_split(dataset, [55000, 5000])
 
-    model_path = download_model(name="jirka/kaggle/lit-auto-encoder-simple:latest", download_dir="my_models")
+    model_path = download_model(name=MY_MODEL_NAME, download_dir="my_models")
     print(f"model: {model_path}")
     # autoencoder = LitAutoEncoder.load_from_checkpoint(checkpoint_path=model_path)
 

examples/train-simple.py

@@ -5,6 +5,11 @@
 from litmodels import upload_model
 from sample_model import LitAutoEncoder
 
+# Define the model name - this should be unique to your model
+# The format is <organization>/<teamspace>/<model-name>
+MY_MODEL_NAME = "jirka/kaggle/lit-auto-encoder-callback"
+
+
 if __name__ == "__main__":
     dataset = tv.datasets.MNIST(".", download=True, transform=tv.transforms.ToTensor())
     train, val = data.random_split(dataset, [55000, 5000])
@@ -30,4 +35,4 @@
         data.DataLoader(val, batch_size=256),
     )
     print(f"last: {vars(checkpoint_callback)}")
-    upload_model(path=checkpoint_callback.last_model_path, name="jirka/kaggle/lit-auto-encoder-simple")
+    upload_model(path=checkpoint_callback.last_model_path, name=MY_MODEL_NAME)