diff --git a/generate_sweep_params.py b/generate_sweep_params.py index c0bad0a..38af571 100644 --- a/generate_sweep_params.py +++ b/generate_sweep_params.py @@ -2,6 +2,51 @@ from trainer.utils.json_stuff import save_as_json import itertools import copy import os +import random +random.seed(0) + +GPU_IDS = [1,2,3] + +def divide_list(lst, n): + """ + Divide a list into N equal parts. + + Parameters: + lst (list): The list to be divided. + n (int): The number of parts to divide the list into. + + Returns: + list of lists: A list containing N sublists, each of which is a part of the original list. + """ + if n <= 0: + raise ValueError("Number of parts must be greater than 0.") + if n > len(lst): + raise ValueError("Number of parts cannot be greater than the length of the list.") + + # Calculate the size of each part + k, m = divmod(len(lst), n) + + # Create the divided parts + return [lst[i * k + min(i, m):(i + 1) * k + min(i + 1, m)] for i in range(n)] + +def generate_sh_file(commands, filename="script.sh"): + """ + Generates a .sh file with each command from the list written on a new line. + + :param commands: List of commands to be written to the .sh file. + :param filename: Name of the .sh file to be created. Default is 'script.sh'. + """ + with open(filename, 'w') as file: + for command in commands: + file.write(command + '\n') + print(f"Saved: {filename}") + +run_commands_dir = f"./sd3_sweep_commands" + +os.system( + f"rm -rf {run_commands_dir} && mkdir -p {run_commands_dir}" +) + config_folder = "./sd3_face_sweep_configs" os.system(f"rm -rf {config_folder}") @@ -74,6 +119,7 @@ default_config = { keys, values = zip(*sweep_params.items()) combinations = [dict(zip(keys, combination)) for combination in itertools.product(*values)] +all_config_paths = [] for index, c in enumerate(combinations): config = copy.deepcopy(default_config) filename = f"{index}" @@ -101,6 +147,32 @@ for index, c in enumerate(combinations): dictionary_or_list=config, filename = config_path ) + all_config_paths.append(config_path) print(f"Saved: {config_path}") -print(f"Total: {index+1} configs") \ No newline at end of file +print(f"Total: {index+1} configs") + +all_commands = [] + +for c in all_config_paths: + command = f"python3 main_sd3.py {c}" + all_commands.append(command) + +random.shuffle(all_commands) + +all_commands_split_by_gpu = divide_list( + lst = all_commands, + n = len(GPU_IDS) +) + +for index, gpu_id in enumerate(GPU_IDS): + commands_on_single_gpu = [ + f"CUDA_VISIBLE_DEVICES={gpu_id} {x}" for x in all_commands_split_by_gpu[index] + ] + generate_sh_file( + commands = commands_on_single_gpu, + filename = os.path.join( + run_commands_dir, + f"run_on_gpu_{gpu_id}.sh" + ) + ) \ No newline at end of file