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.
40 lines
1.1 KiB
40 lines
1.1 KiB
from random import sample
|
|
|
|
from lightning_sdk import Machine, Studio
|
|
|
|
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]
|
|
learning_rate_values = [1e-6, 1e-5, 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, lr, bs, me)
|
|
for alpha in alpha_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
|
|
search_params = sample(all_params, NUM_JOBS)
|
|
|
|
# start all jobs on an A10G GPU with names containing an index
|
|
for idx, (a, lr, bs, me) in enumerate(search_params):
|
|
cmd = f"python main.py --alpha {a} --lr {lr} --bs {bs} --max_epochs {me}"
|
|
job_name = f"color-exp-{idx}"
|
|
job_plugin.run(cmd, machine=Machine.T4, name=job_name)
|
|
|