import neurenix as nx
from neurenix.automl import BayesianOptimization, ENAS
# Stage 1: Find good architecture
print("Stage 1: Architecture Search")
arch_search_space = {
'num_layers': [2, 3, 4],
'hidden_size': [64, 128, 256],
'kernel_size': [3, 5, 7]
}
enas = ENAS(search_space=arch_search_space, max_trials=20)
best_architecture = enas.search(lambda m: quick_train(m, epochs=3))
# Stage 2: Optimize hyperparameters for best architecture
print("Stage 2: Hyperparameter Optimization")
hparam_space = {
'learning_rate': [0.0001, 0.0003, 0.001, 0.003, 0.01],
'batch_size': [16, 32, 64],
'weight_decay': [0, 1e-5, 1e-4, 1e-3],
'optimizer': ['adam', 'adamw', 'sgd']
}
def objective(params):
model = best_architecture.clone()
if params['optimizer'] == 'adam':
optimizer = nx.optim.Adam(
model.parameters(),
lr=params['learning_rate'],
weight_decay=params['weight_decay']
)
elif params['optimizer'] == 'adamw':
optimizer = nx.optim.AdamW(
model.parameters(),
lr=params['learning_rate'],
weight_decay=params['weight_decay']
)
else:
optimizer = nx.optim.SGD(
model.parameters(),
lr=params['learning_rate'],
weight_decay=params['weight_decay'],
momentum=0.9
)
accuracy = train_and_evaluate(model, optimizer, params['batch_size'])
return accuracy
bayes_opt = BayesianOptimization(param_space=hparam_space, max_trials=30)
best_hparams = bayes_opt.search(objective)
# Stage 3: Final training
print("Stage 3: Final Training")
final_model = best_architecture.clone()
final_optimizer = create_optimizer(final_model, best_hparams)
train_full(final_model, final_optimizer, epochs=100)
print(f"Final test accuracy: {evaluate(final_model, test_loader):.2f}%")