code update

This commit is contained in:
Alex Blank
2025-05-20 23:39:23 +02:00
parent cb7896900b
commit d8b9ccfe99
4 changed files with 76 additions and 45 deletions
+2 -1
View File
@@ -60,6 +60,7 @@ run_configuration = {
}, },
}, },
{ {
"batch_size": 64,
"model_parameters": { "model_parameters": {
"cnn_channels": 64, "cnn_channels": 64,
"kernel_size": 3, "kernel_size": 3,
@@ -69,7 +70,7 @@ run_configuration = {
}, },
}, },
{ {
"batch_size": 128, "batch_size": 64,
"model_parameters": { "model_parameters": {
"cnn_channels": 128, "cnn_channels": 128,
"kernel_size": 5, "kernel_size": 5,
+55 -13
View File
@@ -1,8 +1,8 @@
#!/bin/bash #!/bin/bash
# Parameter check # Parameter check
if [[ $# -ne 5 ]]; then if [[ $# -ne 7 ]]; then
echo "Usage: $0 <run_config_module> <partition> <gpu_type> <num_gpus> <num_cpus_per_gpu>" echo "Usage: $0 <run_config_module> <partition> <gpu_type> <num_gpus> <num_cpus_per_gpu> <item_limit> <time_limit>"
exit 1 exit 1
fi fi
@@ -12,10 +12,47 @@ partition=$2
gpu_type=$3 gpu_type=$3
num_gpus=$4 num_gpus=$4
num_cpus_per_gpu=$5 num_cpus_per_gpu=$5
item_limit=$6
time_limit=$7
# Validate parameters
if [[ ! -f "$run_config" ]]; then
echo "Error: Run configuration file '$run_config' does not exist."
exit 1
fi
if [[ ! "$partition" =~ ^(clara|paula)$ ]]; then
echo "Error: Invalid partition '$partition'. Must be 'clara' or 'paula'."
exit 1
fi
valid_gpus=("a30" "v100" "rtx2080ti")
if [[ ! " ${valid_gpus[@]} " =~ " $gpu_type " ]]; then
echo "Error: Invalid GPU type '$gpu_type'. Must be one of: ${valid_gpus[*]}."
exit 1
fi
if ! [[ "$num_gpus" =~ ^[0-9]+$ ]] || ! [[ "$num_cpus_per_gpu" =~ ^[0-9]+$ ]]; then
echo "Error: Number of GPUs and CPUs per GPU must be positive integers."
exit 1
fi
# item limit must be an integer, either positive or -1
if ! [[ "$item_limit" =~ ^-?[0-9]+$ ]]; then
echo "Error: Item limit must be an integer."
exit 1
fi
if [[ "$item_limit" -lt -1 ]]; then
echo "Error: Item limit must be -1 or a positive integer."
exit 1
fi
# Derive run name # Derive run name
run_name=$(basename "$run_config") run_name=$(basename "$run_config")
run_name="${run_name%.*}" run_name="${run_name%.*}"
ram_per_gpu=16
num_cpus=$(($num_gpus * $num_cpus_per_gpu))
ram=$(($num_gpus * $ram_per_gpu))
# Confirm inputs # Confirm inputs
echo "Starting run with configuration: $run_config" echo "Starting run with configuration: $run_config"
@@ -23,26 +60,26 @@ echo "Partition: $partition"
echo "GPU type: $gpu_type" echo "GPU type: $gpu_type"
echo "Number of GPUs: $num_gpus" echo "Number of GPUs: $num_gpus"
echo "Number of CPUs per GPU: $num_cpus_per_gpu" echo "Number of CPUs per GPU: $num_cpus_per_gpu"
echo "Total CPUs: $num_cpus"
echo "Total RAM: $ram GB"
echo "Item limit: $item_limit"
# Set paths # Set paths
current_dir_path=$(dirname "$(realpath "$0")") current_dir_path=$(dirname "$(realpath "$0")")
venv_path=$(realpath "$current_dir_path/../../venv/bin/activate") venv_path=$(realpath "$current_dir_path/../../venv/bin/activate")
# Create log directory
log_dir="/work/rr41qemu-MA/logs/${run_name}_$(date +%Y-%m-%d_%H-%M-%S)"
mkdir -p "$log_dir"
echo "Logging to $log_dir"
# Submit job # Submit job
sbatch <<EOF sbatch <<EOF
#!/bin/bash #!/bin/bash
#SBATCH --job-name=$run_name #SBATCH --job-name=$run_name
#SBATCH --output=$log_dir/log.out #SBATCH --output=/work/rr41qemu-MA/logs/%j.out
#SBATCH --error=$log_dir/log.err #SBATCH --error=/work/rr41qemu-MA/logs/%j.err
#SBATCH --time=48:00:00 #SBATCH --time=$time_limit
#SBATCH --ntasks=1 #SBATCH --ntasks=1
#SBATCH --cpus-per-task=$(($num_gpus * $num_cpus_per_gpu)) #SBATCH --nodes=1
#SBATCH --mem=32G #SBATCH --ntasks-per-node=1
#SBATCH --cpus-per-task=$num_cpus
#SBATCH --mem=${ram}G
#SBATCH --partition=$partition #SBATCH --partition=$partition
#SBATCH --gpus=$gpu_type:$num_gpus #SBATCH --gpus=$gpu_type:$num_gpus
@@ -52,7 +89,12 @@ source $venv_path
echo "Loading python 3.10..." echo "Loading python 3.10..."
module load Python/3.10.4-GCCcore-11.3.0 module load Python/3.10.4-GCCcore-11.3.0
# create log directory
#log_dir="/work/rr41qemu-MA/logs/${SLURM_JOB_ID}_$(date +%Y-%m-%d_%H-%M-%S)"
#mkdir -p "$log_dir"
#echo "Logging to $log_dir"
cd $current_dir_path cd $current_dir_path
torchrun --nproc_per_node=$num_gpus training_wrapper.py $run_config --item_limit -1 torchrun --nproc_per_node=$num_gpus training_wrapper.py $run_config --item_limit $item_limit
EOF EOF
+4 -8
View File
@@ -182,9 +182,6 @@ if __name__ == "__main__":
raise ValueError( raise ValueError(
"No results directory specified. Please set the RESULTS_ROOT_DIR environment variable or provide a results_dir argument.") "No results directory specified. Please set the RESULTS_ROOT_DIR environment variable or provide a results_dir argument.")
if not os.path.exists(results_dir):
os.makedirs(results_dir)
if not args.lmdb_root_dir: if not args.lmdb_root_dir:
if "base_lmdb_root_dir" not in run_configuration: if "base_lmdb_root_dir" not in run_configuration:
lmdb_root_dir = os.getenv("LMDB_ROOT_DIR") lmdb_root_dir = os.getenv("LMDB_ROOT_DIR")
@@ -197,9 +194,6 @@ if __name__ == "__main__":
raise ValueError( raise ValueError(
"No LMDB root directory specified. Please set the LMDB_ROOT_DIR environment variable or provide a lmdb_root_dir argument.") "No LMDB root directory specified. Please set the LMDB_ROOT_DIR environment variable or provide a lmdb_root_dir argument.")
if not os.path.exists(lmdb_root_dir):
os.makedirs(lmdb_root_dir)
if not args.log_dir: if not args.log_dir:
if "base_log_dir" not in run_configuration: if "base_log_dir" not in run_configuration:
log_dir = os.getenv("LOG_DIR") log_dir = os.getenv("LOG_DIR")
@@ -210,8 +204,6 @@ if __name__ == "__main__":
if log_dir is None: if log_dir is None:
raise ValueError( raise ValueError(
"No log directory specified. Please set the LOG_DIR environment variable or provide a log_dir argument.") "No log directory specified. Please set the LOG_DIR environment variable or provide a log_dir argument.")
if not os.path.exists(log_dir):
os.makedirs(log_dir)
try: try:
item_limit = args.item_limit if args.item_limit else run_configuration["item_limit"] item_limit = args.item_limit if args.item_limit else run_configuration["item_limit"]
@@ -226,6 +218,10 @@ if __name__ == "__main__":
# set up run_id on rank 0 # set up run_id on rank 0
if local_rank == 0: if local_rank == 0:
run_id = f"{run_name}_{datetime.now().strftime('%Y%m%d_%H%M%S')}" run_id = f"{run_name}_{datetime.now().strftime('%Y%m%d_%H%M%S')}"
if not os.path.exists(log_dir):
os.makedirs(log_dir)
if not os.path.exists(results_dir):
os.makedirs(results_dir)
else: else:
run_id = None run_id = None
+15 -23
View File
@@ -132,12 +132,12 @@ def train_model(model: nn.Module,
all_train_ids = train_dataset.lmdb_keys all_train_ids = train_dataset.lmdb_keys
train_subsets = [get_ranked_ids(all_train_ids, i, local_rank, world_size) for i in range(num_epochs)] train_subsets = [get_ranked_ids(all_train_ids, i, local_rank, world_size) for i in range(num_epochs)]
# calc total number of steps for gpu, as it is dependent on subsets # calc total number of steps for gpu, as it is dependent on subsets
local_steps = [train_dataset.get_length_of_data_subset(subset) for subset in train_subsets] local_train_steps = [train_dataset.get_length_of_data_subset(subset) for subset in train_subsets]
total_train_steps = sum(local_steps) total_train_steps = sum(local_train_steps)
all_val_ids = val_dataset.lmdb_keys all_val_ids = val_dataset.lmdb_keys
val_subsets = [get_ranked_ids(all_val_ids, i, local_rank, world_size) for i in range(num_epochs)] val_subsets = [get_ranked_ids(all_val_ids, i, local_rank, world_size) for i in range(num_epochs)]
val_subset_lengths = [len(subset) for subset in val_subsets] local_val_steps = [val_dataset.get_length_of_data_subset(subset) for subset in val_subsets]
def get_synced_values(local_values): def get_synced_values(local_values):
""" """
@@ -156,16 +156,16 @@ def train_model(model: nn.Module,
return gathered_values return gathered_values
# sync lengths of subsets and adjust to minimum length for equal sized training lengths # sync lengths of subsets and adjust to minimum length for equal sized training lengths
global_lengths = get_synced_values(local_steps) global_train_lengths = get_synced_values(local_train_steps)
global_lengths = [x.cpu().numpy() for x in global_lengths] global_train_lengths = [x.cpu().numpy() for x in global_train_lengths]
train_epoch_lengths = list() train_epoch_lengths = list()
for i in range(len(global_lengths[0])): for i in range(len(global_train_lengths[0])):
current_epoch_lengths = [x[i] for x in global_lengths] current_epoch_lengths = [x[i] for x in global_train_lengths]
train_epoch_lengths.append(int(min(current_epoch_lengths))) train_epoch_lengths.append(int(min(current_epoch_lengths)))
print(f"Rank {local_rank}: Global lengths: {global_lengths} cut to {train_epoch_lengths}") print(f"Rank {local_rank}: Global lengths: {global_train_lengths} cut to {train_epoch_lengths}")
# also sync the val subsets # also sync the val subsets
global_val_lengths = get_synced_values(val_subset_lengths) global_val_lengths = get_synced_values(local_val_steps)
global_val_lengths = [x.cpu().numpy() for x in global_val_lengths] global_val_lengths = [x.cpu().numpy() for x in global_val_lengths]
val_epoch_lengths = list() val_epoch_lengths = list()
for i in range(len(global_val_lengths[0])): for i in range(len(global_val_lengths[0])):
@@ -249,6 +249,7 @@ def train_model(model: nn.Module,
num_workers=num_dataloader_workers, num_workers=num_dataloader_workers,
# set multiprocessing start method to spawn # set multiprocessing start method to spawn
# multiprocessing_context="forkserver", # multiprocessing_context="forkserver",
multiprocessing_context="spawn",
) )
train_length = train_epoch_lengths[epoch - 1] train_length = train_epoch_lengths[epoch - 1]
val_dataloader = DataLoader( val_dataloader = DataLoader(
@@ -257,6 +258,7 @@ def train_model(model: nn.Module,
num_workers=num_dataloader_workers, num_workers=num_dataloader_workers,
# set multiprocessing start method to spawn # set multiprocessing start method to spawn
# multiprocessing_context="forkserver", # multiprocessing_context="forkserver",
multiprocessing_context="spawn",
) )
val_length = val_epoch_lengths[epoch - 1] val_length = val_epoch_lengths[epoch - 1]
@@ -270,13 +272,8 @@ def train_model(model: nn.Module,
model_configuration) model_configuration)
except StopIteration: except StopIteration:
# if the iterator is exhausted, reset it # if the iterator is exhausted, reset it
logger.info(f"Rank {local_rank}: Iterator exhausted, resetting it.") logger.info(f"Rank {local_rank}: Iterator exhausted, continuing to next epoch.")
iterator = iter(train_dataloader) break
loss = batch_loss_fn(model,
iterator,
loss_functions,
device,
model_configuration)
optimizer.zero_grad() optimizer.zero_grad()
loss.backward() loss.backward()
@@ -317,13 +314,8 @@ def train_model(model: nn.Module,
model_configuration) model_configuration)
except StopIteration: except StopIteration:
# if the iterator is exhausted, reset it # if the iterator is exhausted, reset it
logger.info(f"Rank {local_rank}: Iterator exhausted, resetting it.") logger.info(f"Rank {local_rank}: Iterator exhausted, continuing to next epoch.")
val_iter = iter(val_dataloader) break
loss = batch_loss_fn(model,
val_iter,
loss_functions,
device,
model_configuration)
total_val_loss += loss.item() total_val_loss += loss.item()