You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

39 lines
1.4 KiB

11 months ago
from lightning_sdk import Studio, Machine
from random import sample
NUM_JOBS = 4
# reference to the current studio
# if you run outside of Lightning, you can pass the Studio name
studio = Studio()
# use the jobs plugin
studio.install_plugin('jobs')
job_plugin = studio.installed_plugins['jobs']
# do a sweep over learning rates
# Define the ranges or sets of values for each hyperparameter
alpha_values = [0.1, 0.3, 0.5, 0.7, 0.9]
distinct_threshold_values = [0.5, 0.6, 0.7, 0.8]
learning_rate_values = [1e-6, 1-e5, 1e-4, 1e-3, 1e-2]
batch_size_values = [32, 64, 128]
max_epochs_values = [10000]
# Generate all possible combinations of hyperparameters
all_params = [(alpha, dt, lr, bs, me) for alpha in alpha_values
for dt in distinct_threshold_values
for lr in learning_rate_values
for bs in batch_size_values
for me in max_epochs_values]
# perform random search with a limit
random_search_params = sample(search_params, NUM_JOBS)
# start all jobs on an A10G GPU with names containing an index
for idx, (a, thresh, lr, bs, max_epochs) in enumerate(search_params):
cmd = f'python main.py --alpha {a} -D {thresh} --lr {lr} --bs {bs} --max_epochs {max_epochs}'
job_name = f'color-exp-{idx}'
job_plugin.run(cmd, machine=Machine.T4, name=job_name)