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": {
"cnn_channels": 64,
"kernel_size": 3,
@@ -69,7 +70,7 @@ run_configuration = {
},
},
{
"batch_size": 128,
"batch_size": 64,
"model_parameters": {
"cnn_channels": 128,
"kernel_size": 5,
+55 -13
View File
@@ -1,8 +1,8 @@
#!/bin/bash
# Parameter check
if [[ $# -ne 5 ]]; then
echo "Usage: $0 <run_config_module> <partition> <gpu_type> <num_gpus> <num_cpus_per_gpu>"
if [[ $# -ne 7 ]]; then
echo "Usage: $0 <run_config_module> <partition> <gpu_type> <num_gpus> <num_cpus_per_gpu> <item_limit> <time_limit>"
exit 1
fi
@@ -12,10 +12,47 @@ partition=$2
gpu_type=$3
num_gpus=$4
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
run_name=$(basename "$run_config")
run_name="${run_name%.*}"
ram_per_gpu=16
num_cpus=$(($num_gpus * $num_cpus_per_gpu))
ram=$(($num_gpus * $ram_per_gpu))
# Confirm inputs
echo "Starting run with configuration: $run_config"
@@ -23,26 +60,26 @@ echo "Partition: $partition"
echo "GPU type: $gpu_type"
echo "Number of GPUs: $num_gpus"
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
current_dir_path=$(dirname "$(realpath "$0")")
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
sbatch <<EOF
#!/bin/bash
#SBATCH --job-name=$run_name
#SBATCH --output=$log_dir/log.out
#SBATCH --error=$log_dir/log.err
#SBATCH --time=48:00:00
#SBATCH --output=/work/rr41qemu-MA/logs/%j.out
#SBATCH --error=/work/rr41qemu-MA/logs/%j.err
#SBATCH --time=$time_limit
#SBATCH --ntasks=1
#SBATCH --cpus-per-task=$(($num_gpus * $num_cpus_per_gpu))
#SBATCH --mem=32G
#SBATCH --nodes=1
#SBATCH --ntasks-per-node=1
#SBATCH --cpus-per-task=$num_cpus
#SBATCH --mem=${ram}G
#SBATCH --partition=$partition
#SBATCH --gpus=$gpu_type:$num_gpus
@@ -52,7 +89,12 @@ source $venv_path
echo "Loading python 3.10..."
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
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
+4 -8
View File
@@ -182,9 +182,6 @@ if __name__ == "__main__":
raise ValueError(
"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 "base_lmdb_root_dir" not in run_configuration:
lmdb_root_dir = os.getenv("LMDB_ROOT_DIR")
@@ -197,9 +194,6 @@ if __name__ == "__main__":
raise ValueError(
"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 "base_log_dir" not in run_configuration:
log_dir = os.getenv("LOG_DIR")
@@ -210,8 +204,6 @@ if __name__ == "__main__":
if log_dir is None:
raise ValueError(
"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:
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
if local_rank == 0:
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:
run_id = None
+15 -23
View File
@@ -132,12 +132,12 @@ def train_model(model: nn.Module,
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)]
# 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]
total_train_steps = sum(local_steps)
local_train_steps = [train_dataset.get_length_of_data_subset(subset) for subset in train_subsets]
total_train_steps = sum(local_train_steps)
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_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):
"""
@@ -156,16 +156,16 @@ def train_model(model: nn.Module,
return gathered_values
# sync lengths of subsets and adjust to minimum length for equal sized training lengths
global_lengths = get_synced_values(local_steps)
global_lengths = [x.cpu().numpy() for x in global_lengths]
global_train_lengths = get_synced_values(local_train_steps)
global_train_lengths = [x.cpu().numpy() for x in global_train_lengths]
train_epoch_lengths = list()
for i in range(len(global_lengths[0])):
current_epoch_lengths = [x[i] for x in global_lengths]
for i in range(len(global_train_lengths[0])):
current_epoch_lengths = [x[i] for x in global_train_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
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]
val_epoch_lengths = list()
for i in range(len(global_val_lengths[0])):
@@ -249,6 +249,7 @@ def train_model(model: nn.Module,
num_workers=num_dataloader_workers,
# set multiprocessing start method to spawn
# multiprocessing_context="forkserver",
multiprocessing_context="spawn",
)
train_length = train_epoch_lengths[epoch - 1]
val_dataloader = DataLoader(
@@ -257,6 +258,7 @@ def train_model(model: nn.Module,
num_workers=num_dataloader_workers,
# set multiprocessing start method to spawn
# multiprocessing_context="forkserver",
multiprocessing_context="spawn",
)
val_length = val_epoch_lengths[epoch - 1]
@@ -270,13 +272,8 @@ def train_model(model: nn.Module,
model_configuration)
except StopIteration:
# if the iterator is exhausted, reset it
logger.info(f"Rank {local_rank}: Iterator exhausted, resetting it.")
iterator = iter(train_dataloader)
loss = batch_loss_fn(model,
iterator,
loss_functions,
device,
model_configuration)
logger.info(f"Rank {local_rank}: Iterator exhausted, continuing to next epoch.")
break
optimizer.zero_grad()
loss.backward()
@@ -317,13 +314,8 @@ def train_model(model: nn.Module,
model_configuration)
except StopIteration:
# if the iterator is exhausted, reset it
logger.info(f"Rank {local_rank}: Iterator exhausted, resetting it.")
val_iter = iter(val_dataloader)
loss = batch_loss_fn(model,
val_iter,
loss_functions,
device,
model_configuration)
logger.info(f"Rank {local_rank}: Iterator exhausted, continuing to next epoch.")
break
total_val_loss += loss.item()