save different sh files for each gpu

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")
+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"
+            )
+        )