code cleanup

main.py

@@ -0,0 +1,64 @@
+from argparse import ArgumentParser
+from math import ceil
+
+import torch
+from torch.utils.data import DataLoader
+
+from dataset_loaders import EnWik9DataSet, LoremIpsumDataset, Dataset
+from trainers import OptunaTrainer, Trainer, FullTrainer
+
+BATCH_SIZE = 64
+
+if torch.cuda.is_available():
+    DEVICE = "cuda"
+elif torch.backends.mps.is_available():
+    DEVICE = "mps"
+else:
+    DEVICE = "cpu"
+
+# hyper parameters
+context_length = 128
+
+if __name__ == "__main__":
+    print(f"Running on device: {DEVICE}...")
+    parser = ArgumentParser()
+    parser.add_argument("--method", choices=["optuna", "train"], required=True)
+    parser.add_argument("--models-path", type=str, required=False)
+    args = parser.parse_args()
+
+    print("Loading in the dataset...")
+    if args.method == "train":
+        dataset: Dataset = EnWik9DataSet(transform=lambda x: x.to(DEVICE))
+    elif args.method == "optuna":
+        dataset: Dataset = LoremIpsumDataset(transform=lambda x: x.to(DEVICE))
+    else:
+        raise ValueError(f"Unknown method: {args.method}")
+
+    dataset_length = len(dataset)
+    print(f"Dataset size = {dataset_length}")
+
+    training_size = ceil(0.8 * dataset_length)
+
+    print(f"Training set size = {training_size}, Validation set size {dataset_length - training_size}")
+
+    train_set, validate_set = torch.utils.data.random_split(dataset,
+                                                            [training_size, dataset_length - training_size])
+    training_loader = DataLoader(train_set, batch_size=BATCH_SIZE, shuffle=True)
+    validation_loader = DataLoader(validate_set, batch_size=BATCH_SIZE, shuffle=False)
+    loss_fn = torch.nn.CrossEntropyLoss()
+
+    model = None
+    if args.model_path is not None:
+        print("Loading the models...")
+        model = torch.load(args.model_path)
+
+    trainer: Trainer = OptunaTrainer() if args.method == "optuna" else FullTrainer()
+
+    trainer.execute(
+        model=model,
+        train_loader=training_loader,
+        validation_loader=validation_loader,
+        loss_fn=loss_fn,
+        n_epochs=200,
+        device=DEVICE
+    )