device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print('Using device:', device)
X_train_t = torch.tensor(X_train_proc, dtype=torch.float32)
y_train_t = torch.tensor(y_train.values.reshape(-1, 1), dtype=torch.float32)
X_val_t = torch.tensor(X_val_proc, dtype=torch.float32)
y_val_t = torch.tensor(y_val.values.reshape(-1, 1), dtype=torch.float32)
X_test_t = torch.tensor(X_test_proc, dtype=torch.float32)
y_test_t = torch.tensor(y_test.values.reshape(-1, 1), dtype=torch.float32)
train_dataset = TensorDataset(X_train_t, y_train_t)
val_dataset = TensorDataset(X_val_t, y_val_t)
class NeuralRegressor(nn.Module):
def __init__(self, in_features, hidden_layers=(128, 64), dropout=0.15):
super().__init__()
layers = []
prev = in_features
for h in hidden_layers:
layers += [nn.Linear(prev, h), nn.ReLU(), nn.BatchNorm1d(h), nn.Dropout(dropout)]
prev = h
layers += [nn.Linear(prev, 1)]
self.net = nn.Sequential(*layers)
def forward(self, x):
return self.net(x)
def train_one_config(config, seed=SEED):
torch.manual_seed(seed)
np.random.seed(seed)
random.seed(seed)
model = NeuralRegressor(
in_features=X_train_proc.shape[1],
hidden_layers=config['hidden_layers'],
dropout=config['dropout']
).to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=config['lr'], weight_decay=config['weight_decay'])
criterion = nn.MSELoss()
train_loader = DataLoader(train_dataset, batch_size=config['batch_size'], shuffle=True)
val_loader = DataLoader(val_dataset, batch_size=4096, shuffle=False)
best_state = None
best_val_rmse = float('inf')
patience = config['patience']
wait = 0
history = []
for epoch in range(config['epochs']):
model.train()
train_losses = []
for xb, yb in train_loader:
xb = xb.to(device)
yb = yb.to(device)
optimizer.zero_grad()
pred = model(xb)
loss = criterion(pred, yb)
loss.backward()
optimizer.step()
train_losses.append(loss.item())
model.eval()
val_preds = []
val_true = []
with torch.no_grad():
for xb, yb in val_loader:
xb = xb.to(device)
out = model(xb).cpu().numpy().ravel()
val_preds.append(out)
val_true.append(yb.numpy().ravel())
yv_pred = np.concatenate(val_preds)
yv_true = np.concatenate(val_true)
val_rmse = np.sqrt(mean_squared_error(yv_true, yv_pred))
history.append((epoch + 1, float(np.mean(train_losses)), float(val_rmse)))
if val_rmse < best_val_rmse:
best_val_rmse = val_rmse
best_state = {k: v.cpu().clone() for k, v in model.state_dict().items()}
wait = 0
else:
wait += 1
if wait >= patience:
break
model.load_state_dict(best_state)
return model, best_val_rmse, pd.DataFrame(history, columns=['epoch', 'train_mse', 'val_rmse'])
nn_grid = [
{'hidden_layers': (128, 64), 'dropout': 0.10, 'lr': 1e-3, 'weight_decay': 1e-5, 'batch_size': 1024, 'epochs': 45, 'patience': 6},
{'hidden_layers': (256, 128, 64), 'dropout': 0.15, 'lr': 7e-4, 'weight_decay': 1e-5, 'batch_size': 1024, 'epochs': 55, 'patience': 7},
{'hidden_layers': (192, 96), 'dropout': 0.20, 'lr': 6e-4, 'weight_decay': 5e-5, 'batch_size': 2048, 'epochs': 60, 'patience': 7}
]
nn_trials = []
best_nn_model = None
best_nn_history = None
best_nn_config = None
best_nn_val_rmse = float('inf')
for config in nn_grid:
model, val_rmse, history_df = train_one_config(config, seed=SEED)
nn_trials.append({
'config': str(config),
'val_rmse': val_rmse
})
if val_rmse < best_nn_val_rmse:
best_nn_val_rmse = val_rmse
best_nn_model = model
best_nn_history = history_df
best_nn_config = config
nn_trials_df = pd.DataFrame(nn_trials).sort_values('val_rmse')
display(nn_trials_df)
print('Best NN config:', best_nn_config)
print('Best validation RMSE:', round(best_nn_val_rmse, 4))
