MethSurvPredictor / hyperparameters.py
csycsycsy's picture
Upload 10 files
84db192 verified
import optuna
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, TensorDataset
import pandas as pd
import numpy as np
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler
import matplotlib.pyplot as plt
import os
script_path = os.path.abspath(__file__)
script_dir = os.path.dirname(script_path)
os.chdir(script_dir)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Using device: {device}")
data = pd.read_csv('data.csv')
X = data.drop(columns=['OS.time']).values
y = data['OS.time'].values
scaler = StandardScaler()
X_scaled = scaler.fit_transform(X)
X_train, X_test, y_train, y_test = train_test_split(X_scaled, y, test_size=0.2, random_state=42)
X_train_tensor = torch.tensor(X_train, dtype=torch.float32).to(device)
y_train_tensor = torch.tensor(y_train, dtype=torch.float32).view(-1, 1).to(device)
X_test_tensor = torch.tensor(X_test, dtype=torch.float32).to(device)
y_test_tensor = torch.tensor(y_test, dtype=torch.float32).view(-1, 1).to(device)
train_dataset = TensorDataset(X_train_tensor, y_train_tensor)
test_dataset = TensorDataset(X_test_tensor, y_test_tensor)
train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=32, shuffle=False)
class SimpleNN(nn.Module):
def __init__(self, input_dim, hidden_dim, num_layers, dropout_rate):
super(SimpleNN, self).__init__()
self.layers = nn.ModuleList()
last_dim = input_dim
for _ in range(num_layers):
self.layers.append(nn.Linear(last_dim, hidden_dim))
self.layers.append(nn.ReLU())
self.layers.append(nn.Dropout(dropout_rate))
last_dim = hidden_dim
self.layers.append(nn.Linear(last_dim, 1))
def forward(self, x):
for layer in self.layers:
x = layer(x)
return x
def weights_init(m):
if isinstance(m, nn.Linear):
nn.init.kaiming_uniform_(m.weight)
nn.init.zeros_(m.bias)
def objective(trial):
num_layers = trial.suggest_int('num_layers', 2, 5)
hidden_dim = trial.suggest_int('hidden_dim', 50, 200)
dropout_rate = trial.suggest_float('dropout_rate', 0.2, 0.5)
momentum = trial.suggest_float('momentum', 0.5, 0.9)
num_epochs = trial.suggest_int('num_epochs', 6000, 10000)
model = SimpleNN(X_train.shape[1], hidden_dim, num_layers, dropout_rate).to(device)
model.apply(weights_init)
criterion = nn.MSELoss()
optimizer = optim.SGD(model.parameters(), lr=0.0001, momentum=momentum)
for epoch in range(num_epochs):
model.train()
for inputs, targets in train_loader:
optimizer.zero_grad()
outputs = model(inputs)
loss = criterion(outputs, targets)
loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
optimizer.step()
model.eval()
test_loss = 0
with torch.no_grad():
for inputs, targets in test_loader:
outputs = model(inputs)
test_loss += criterion(outputs, targets).item() * inputs.size(0)
test_loss /= len(test_loader.dataset)
trial.report(test_loss, epoch)
if trial.should_prune():
raise optuna.exceptions.TrialPruned()
return test_loss
study = optuna.create_study(direction='minimize')
study.optimize(objective, n_trials=200)
print(f"Best trial parameters: {study.best_trial.params}")
print(f"Best trial test loss: {study.best_trial.value}")
import optuna.visualization as vis
vis.plot_param_importances(study).show()
vis.plot_parallel_coordinate(study).show()
best_params = study.best_trial.params
model = SimpleNN(X_train.shape[1], best_params['hidden_dim'], best_params['num_layers'], best_params['dropout_rate']).to(device)
model.apply(weights_init)
criterion = nn.MSELoss()
optimizer = optim.SGD(model.parameters(), lr=0.0001, momentum=best_params['momentum'])
test_losses = []
for epoch in range(best_params['num_epochs']):
model.train()
for inputs, targets in train_loader:
optimizer.zero_grad()
outputs = model(inputs)
loss = criterion(outputs, targets)
loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
optimizer.step()
model.eval()
test_loss = 0
with torch.no_grad():
for inputs, targets in test_loader:
outputs = model(inputs)
test_loss += criterion(outputs, targets).item() * inputs.size(0)
test_loss /= len(test_loader.dataset)
if epoch % 100 == 0:
test_losses.append(test_loss)
print(f'Epoch {epoch+1}, Test Loss: {test_loss}')
print("Training completed with best hyperparameters.")
plt.figure(figsize=(10, 5))
plt.plot(range(1, len(test_losses) * 100, 100), test_losses, label='Test Loss')
plt.xlabel('Epoch')
plt.ylabel('Test Loss')
plt.title('Test Loss over Epochs')
plt.legend()
plt.show()