fix/last-autoencoder-test

minecraft_copilot_ml/__main__.py

@@ -84,12 +84,12 @@ def collate_fn(batch: List[MinecraftSchematicsDatasetItemType]) -> MinecraftSche
 
     model = UNet3d(unique_blocks_dict, unique_counts_coefficients=unique_counts_coefficients)
     csv_logger = CSVLogger(save_dir=path_to_output)
-    model_checkpoint = ModelCheckpoint(path_to_output, monitor="val_loss", save_top_k=1, save_last=True, mode="min")
+    model_checkpoint = ModelCheckpoint(path_to_output, monitor="val_accuracy_on_loss_map", save_top_k=1, save_last=True, mode="max")
     trainer = pl.Trainer(logger=csv_logger, callbacks=model_checkpoint, max_epochs=epochs, log_every_n_steps=1)
     trainer.fit(model, train_schematics_dataloader, val_schematics_dataloader)
 
     # Save the best and last model locally
-    logger.info(f"Best val_loss is: {model_checkpoint.best_model_score}")
+    logger.info(f"Best val_accuracy_on_loss_map is: {model_checkpoint.best_model_score}")
     best_model = UNet3d.load_from_checkpoint(
         model_checkpoint.best_model_path,
         unique_blocks_dict=unique_blocks_dict,

minecraft_copilot_ml/data_loader.py

@@ -256,5 +256,7 @@ def get_working_files_and_unique_blocks_and_counts(
             continue
     unique_blocks_dict = {block: idx for idx, block in enumerate(unique_blocks)}
     unique_counts_coefficients = np.array([unique_counts[block] for block in unique_blocks_dict])
-    unique_counts_coefficients = unique_counts_coefficients.max() / unique_counts_coefficients
+    unique_counts_coefficients = unique_counts_coefficients - unique_counts_coefficients.min()
+    unique_counts_coefficients = unique_counts_coefficients / unique_counts_coefficients.max()
+    unique_counts_coefficients = unique_counts_coefficients + 1
     return unique_blocks_dict, unique_counts_coefficients, loaded_schematic_files

minecraft_copilot_ml/model.py

@@ -6,7 +6,7 @@
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
-
+from torch.optim.lr_scheduler import ReduceLROnPlateau
 from minecraft_copilot_ml.data_loader import MinecraftSchematicsDatasetItemType
 
 
@@ -75,7 +75,9 @@ def step(self, batch: MinecraftSchematicsDatasetItemType, batch_idx: int, mode:
         block_maps, noisy_block_maps, block_map_masks, loss_masks = batch
         pre_processed_block_maps = self.pre_process(block_maps)
         pre_processed_noisy_block_maps = self.pre_process(noisy_block_maps).float().unsqueeze(1)
-        tensor_block_map_masks = torch.from_numpy(block_map_masks).float().to("cuda" if torch.cuda.is_available() else "cpu").long()
+        tensor_block_map_masks = (
+            torch.from_numpy(block_map_masks).float().to("cuda" if torch.cuda.is_available() else "cpu").long()
+        )
         tensor_loss_masks = (
             torch.from_numpy(loss_masks).float().to("cuda" if torch.cuda.is_available() else "cpu").long()
         )
@@ -132,7 +134,14 @@ def validation_step(self, batch: MinecraftSchematicsDatasetItemType, batch_idx:
         return self.step(batch, batch_idx, "val")
 
     def configure_optimizers(self) -> Any:
-        return torch.optim.Adam(self.parameters(), lr=1e-3)
+        optimizer = torch.optim.Adam(self.parameters(), lr=0.1)
+        return {
+            "optimizer": optimizer,
+            "lr_scheduler": {
+                "scheduler": ReduceLROnPlateau(optimizer, mode="max", factor=0.1),
+                "monitor": "val_accuracy_on_loss_map",
+            },
+        }
 
     def on_train_start(self) -> None:
         print(self)