Checkpoint

main.py

@@ -17,23 +17,27 @@ def main():
     match args.mode:
         case 'train':
             train(
-                device = device,
-                dataset = args.dataset,
-                data_root = args.data_root,
-                n_trials = 3 if args.debug else None,
-                size = 2**10 if args.debug else None,
-                method = args.method,
+                device=device,
+                dataset=args.dataset,
+                data_root=args.data_root,
+                n_trials=3 if args.debug else None,
+                size=2 ** 10 if args.debug else None,
+                method=args.method,
                 model_name=args.model,
-                model_path = args.model_load_path,
-                model_out = args.model_save_path
+                model_path=args.model_load_path,
+                model_out=args.model_save_path
             )
 
         case 'compress':
-            compress(args.input_file)
+            compress(device=device,
+                     model_path=args.model_load_path,
+                     input_file=args.input_file,
+                     output_file=args.output_file
+                     )
 
         case _:
             raise NotImplementedError(f"Mode {args.mode} is not implemented yet")
-    
+
     print("Done")
 
 

src/args.py

@@ -19,7 +19,7 @@ def parse_arguments():
                              help="Which model to use")
     modelparser.add_argument("--model-load-path", type=str, required=False,
                              help="Filepath to the model to load")
-    modelparser.add_argument("--model-save-path", type=str, required=True,
+    modelparser.add_argument("--model-save-path", type=str, required=False,
                              help="Filepath to the model to save")
 
     fileparser = ArgumentParser(add_help=False)

src/utils/utils.py

@@ -1,3 +1,5 @@
+from os import path
+
 import torch
 from torch.utils.data import TensorDataset
 import matplotlib.pyplot as plt
@@ -14,7 +16,7 @@ def print_distribution(from_to: tuple[int, int], probabilities: list[float]):
     plt.hist(range(from_to[0], from_to[1]), weights=probabilities)
     plt.show()
 
-def print_losses(train_losses: list[float], validation_losses: list[float], show=False):
+def print_losses(train_losses: list[float], validation_losses: list[float], filename: str | None = None, show=False):
     plt.plot(train_losses, label="Training loss")
     plt.plot(validation_losses, label="Validation loss")
     plt.xlabel("Epoch")
@@ -23,7 +25,12 @@ def print_losses(train_losses: list[float], validation_losses: list[float], show
 
     if show:
         plt.show()
-    plt.savefig("losses.png")
+
+    if filename is None:
+        filename = path.join("results", "losses.png")
+
+    print(f"Saving losses to {filename}...")
+    plt.savefig(filename)
 
 
 def load_data(path: str) -> bytes: