code cleanup

trainers/OptunaTrainer.py

@@ -0,0 +1,57 @@
+from typing import Callable
+
+import optuna
+import optuna.trial as tr
+import torch
+from torch import nn as nn
+from torch.utils.data import DataLoader
+
+from .trainer import Trainer
+from ..models.cnn import CNNPredictor
+from .train import train
+
+
+def create_model(trial: tr.Trial, vocab_size: int = 256):
+    hidden_dim = trial.suggest_int("hidden_dim", 64, 512, log=True)
+    embedding_dim = trial.suggest_int("embed_dim", 64, 512, log=True)
+
+    return CNNPredictor(
+        vocab_size=vocab_size,
+        hidden_dim=hidden_dim,
+        embed_dim=embedding_dim,
+    )
+
+
+def objective_function(
+        trial: tr.Trial,
+        training_loader: DataLoader,
+        validation_loader: DataLoader,
+        loss_fn: Callable[[torch.Tensor, torch.Tensor], torch.Tensor],
+        device: str
+):
+    model = create_model(trial).to(device)
+    _, validation_loss = train(model, training_loader, validation_loader, loss_fn)
+    return min(validation_loss)
+
+
+class OptunaTrainer(Trainer):
+    def execute(
+            self,
+            model: nn.Module | None,
+            train_loader: DataLoader,
+            validation_loader: DataLoader,
+            loss_fn: Callable[[torch.Tensor, torch.Tensor], torch.Tensor],
+            n_epochs: int,
+            device: str
+    ) -> None:
+        study = optuna.create_study(study_name="CNN network", direction="minimize")
+        study.optimize(
+            lambda trial: objective_function(trial, train_loader, validation_loader, loss_fn, device),
+            n_trials=20
+        )
+
+        best_params = study.best_trial.params
+        best_model = CNNPredictor(
+            **best_params
+        )
+        torch.save(best_model, f"saved_models/{model.__class__.__name__}.pt")