Actually save the transformer weights

woodRock · Sep 24, 2024 · 5d766d4 · 5d766d4
1 parent d75e9d3
commit 5d766d4
Show file tree

Hide file tree

Showing 8 changed files with 232 additions and 222 deletions.
diff --git a/code/transformer/figures/decoder_attention_map.png b/code/transformer/figures/decoder_attention_map.png
diff --git a/code/transformer/figures/encoder_attention_map.png b/code/transformer/figures/encoder_attention_map.png
diff --git a/code/transformer/figures/model_accuracy.png b/code/transformer/figures/model_accuracy.png
diff --git a/code/transformer/figures/train_confusion_matrix.png b/code/transformer/figures/train_confusion_matrix.png
diff --git a/code/transformer/figures/validation_confusion_matrix.png b/code/transformer/figures/validation_confusion_matrix.png
diff --git a/code/transformer/logs/results_0.log b/code/transformer/logs/results_0.log
diff --git a/code/transformer/multi-task.py b/code/transformer/multi-task.py
@@ -126,6 +126,9 @@ def main():
             num_epochs=args.epochs, 
             patience=args.early_stopping
         )
+
+        # Save the model to disk.
+        torch.save(model.state_dict(), args.file_path)
 
         # finish measuring how long training took
         endTime = time.time()
@@ -141,4 +144,4 @@ def main():
 
 
 if __name__ == "__main__":
-    main()
+    main()
diff --git a/code/transformer/util.py b/code/transformer/util.py
@@ -102,7 +102,8 @@ def random_augmentation(
     return xs, ys
 
 def load_from_file(
-        path: Iterable = ["~/", "Desktop", "fishy-business", "data", "REIMS_data.xlsx"]
+        # path: Iterable = ["~/", "Desktop", "fishy-business", "data", "REIMS_data.xlsx"]
+        path: Iterable = ["/vol","ecrg-solar","woodj4","fishy-business","data", "REIMS_data.xlsx"]
     ) -> pd.DataFrame:
     """ Load the dataset from a file path.
 
@@ -351,4 +352,4 @@ def preprocess_dataset(
         is_data_augmentation=is_data_augmentation,
         batch_size=batch_size
     )
-    return train_loader, val_loader, train_steps, val_steps, data
+    return train_loader, val_loader, train_steps, val_steps, data