chore: Restructure

trainers/OptunaTrainer.py

@@ -1,62 +0,0 @@
-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 __init__(self, n_trials: int | None = None):
-        super().__init__()
-        self.n_trials = n_trials if n_trials is not None else 20
-        print(f"Creating Optuna trainer(n_trials = {self.n_trials})")
-
-    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=self.n_trials
-        )
-
-        best_params = study.best_trial.params
-        best_model = CNNPredictor(
-            **best_params
-        )
-        torch.save(best_model, f"saved_models/{model.__class__.__name__}.pt")