93 lines
No EOL
3.2 KiB
Python
93 lines
No EOL
3.2 KiB
Python
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
import optuna.trial as tr
|
|
from torch.utils.data import DataLoader
|
|
from tqdm import tqdm
|
|
import argparse
|
|
|
|
from optuna_trial import create_model
|
|
from utils import make_context_pairs, load_data
|
|
import optuna
|
|
|
|
# hyper parameters
|
|
context_length = 128
|
|
|
|
def train_and_eval(
|
|
model: nn.Module,
|
|
training_data: bytes,
|
|
validation_data: bytes,
|
|
batch_size: int,
|
|
epochs: int = 100,
|
|
learning_rate: float = 1e-3,
|
|
device: torch.device = torch.device("cpu")
|
|
) -> dict:
|
|
model.to(device)
|
|
optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate)
|
|
training_loader = DataLoader(make_context_pairs(training_data, context_length=context_length), batch_size=batch_size)
|
|
validation_loader= DataLoader(make_context_pairs(validation_data, context_length=context_length), batch_size=batch_size)
|
|
|
|
training_losses = []
|
|
validation_losses = []
|
|
best_val_loss = float("inf")
|
|
|
|
for epoch in range(epochs):
|
|
model.train()
|
|
train_loss = 0
|
|
for x, y in tqdm(training_loader, desc=f"Epoch {epoch}"):
|
|
x, y = x.to(device), y.to(device)
|
|
prediction = model(x)
|
|
loss = F.cross_entropy(prediction, y)
|
|
train_loss += loss.item()
|
|
|
|
optimizer.zero_grad()
|
|
loss.backward()
|
|
optimizer.step()
|
|
training_losses.append(train_loss / len(training_loader))
|
|
|
|
model.eval()
|
|
|
|
with torch.no_grad():
|
|
val_loss = 0
|
|
for x, y in validation_loader:
|
|
x, y = x.to(device), y.to(device)
|
|
prediction = model(x)
|
|
loss = F.cross_entropy(prediction, y)
|
|
val_loss += loss.item()
|
|
validation_losses.append(val_loss / len(validation_loader))
|
|
if validation_losses[-1] < best_val_loss:
|
|
best_val_loss = validation_losses[-1]
|
|
|
|
return {
|
|
"training_losses": training_losses,
|
|
"validation_losses": validation_losses,
|
|
"best_validation_loss": best_val_loss
|
|
}
|
|
|
|
|
|
|
|
def objective_function(trial: tr.Trial, train_data: bytes, validation_data: bytes, batch_size: int):
|
|
model = create_model(trial)
|
|
result = train_and_eval(model, train_data, validation_data, batch_size)
|
|
return result["best_validation_loss"]
|
|
|
|
if __name__ == "__main__":
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--train-data", type=str, required=True)
|
|
parser.add_argument("--validation-data", type=str, required=True)
|
|
parser.add_argument("--batch-size", type=int, default=128)
|
|
|
|
args = parser.parse_args()
|
|
print(args)
|
|
|
|
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
|
train_data = load_data(args.train_data)
|
|
validation_data = load_data(args.validation_data)
|
|
batch_size = args.batch_size
|
|
|
|
print(f"training data length: {len(train_data)}")
|
|
print(f"validation data length: {len(validation_data)}")
|
|
print(f"batch size: {batch_size}")
|
|
|
|
study = optuna.create_study(study_name="CNN network",direction="minimize")
|
|
study.optimize(lambda trial: objective_function(trial, train_data, validation_data, batch_size), n_trials=10) |