fix: one cycle lr

minecraft_copilot_ml/__main__.py

@@ -118,7 +118,11 @@ def collate_fn(batch: List[Tuple[np.ndarray, np.ndarray, np.ndarray]]) -> Tuple[
         val_schematics_dataset, batch_size=batch_size, collate_fn=collate_fn
     )
 
-    model = UNet3d(unique_blocks_dict, unique_counts_coefficients=unique_counts_coefficients)
+    model = UNet3d(
+        unique_blocks_dict,
+        unique_counts_coefficients=unique_counts_coefficients,
+        train_len_dataloader=len(train_schematics_dataloader),
+    )
     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")
     trainer = pl.Trainer(logger=csv_logger, callbacks=model_checkpoint, max_epochs=epochs, log_every_n_steps=1)
@@ -129,12 +133,14 @@ def collate_fn(batch: List[Tuple[np.ndarray, np.ndarray, np.ndarray]]) -> Tuple[
     best_model = UNet3d.load_from_checkpoint(
         model_checkpoint.best_model_path,
         unique_blocks_dict=unique_blocks_dict,
+        train_len_dataloader=len(train_schematics_dataloader),
         unique_counts_coefficients=unique_counts_coefficients,
     )
     torch.save(best_model, os.path.join(path_to_output, "best_model.pth"))
     last_model = UNet3d.load_from_checkpoint(
         model_checkpoint.last_model_path,
         unique_blocks_dict=unique_blocks_dict,
+        train_len_dataloader=len(train_schematics_dataloader),
         unique_counts_coefficients=unique_counts_coefficients,
     )
     torch.save(last_model, os.path.join(path_to_output, "last_model.pth"))

minecraft_copilot_ml/model.py

@@ -2,7 +2,7 @@
 from typing import Any, Dict, Optional, Tuple
 
 import numpy as np
-import pytorch_lightning as pl
+import lightning as pl
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
@@ -23,11 +23,13 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
 class UNet3d(pl.LightningModule):
     def __init__(
         self,
+        train_len_dataloader: int,
         unique_blocks_dict: Dict[str, int],
-        unique_counts_coefficients: Optional[np.ndarray] = None,
         latent_dim: int = 64,
+        unique_counts_coefficients: Optional[np.ndarray] = None,
     ):
         super(UNet3d, self).__init__()
+        self.train_len_dataloader = train_len_dataloader
         self.unique_blocks_dict = unique_blocks_dict
         self.reverse_unique_blocks_dict = {v: k for k, v in unique_blocks_dict.items()}
         self.latent_dim = latent_dim
@@ -127,7 +129,9 @@ def validation_step(self, batch: Tuple[np.ndarray, np.ndarray, np.ndarray], batc
         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=1e-3)
+        scheduler = torch.optim.lr_scheduler.OneCycleLR(optimizer, max_lr=1e-3, total_steps=self.train_len_dataloader)
+        return [optimizer], [scheduler]
 
     def on_train_start(self) -> None:
         print(self)