64 lines
2 KiB
Python
64 lines
2 KiB
Python
from argparse import ArgumentParser
|
|
from math import ceil
|
|
|
|
import torch
|
|
from torch.utils.data import DataLoader
|
|
|
|
from dataset_loaders import dataset_called
|
|
from trainers import OptunaTrainer, Trainer, FullTrainer
|
|
|
|
BATCH_SIZE = 64
|
|
|
|
if torch.accelerator.is_available():
|
|
DEVICE = torch.accelerator.current_accelerator().type
|
|
else:
|
|
DEVICE = "cpu"
|
|
|
|
# hyper parameters
|
|
context_length = 128
|
|
|
|
print(f"Running on device: {DEVICE}...")
|
|
parser = ArgumentParser()
|
|
parser.add_argument("--method", choices=["optuna", "train"], required=True)
|
|
parser.add_argument("--model-path", type=str, required=False)
|
|
|
|
parser.add_argument_group("Data", "Data files or dataset to use")
|
|
parser.add_argument("--data-root", type=str, required=False)
|
|
parser.add_argument("dataset")
|
|
args = parser.parse_args()
|
|
|
|
print("Loading in the dataset...")
|
|
if args.dataset in dataset_called:
|
|
dataset = dataset_called[args.dataset](root=args.data_root, transform=lambda x: x.to(DEVICE))
|
|
else:
|
|
# TODO Allow to import arbitrary files
|
|
raise NotImplementedError(f"Importing external datasets is not implemented yet")
|
|
|
|
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 model...")
|
|
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
|
|
)
|