alex/temperature-based-fertility-prediction-thesis
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

code update

Browse Source
This commit is contained in:
Alex Blank
2025-05-19 22:54:53 +02:00
parent 58305effdb
commit cb7896900b
7 changed files with 238 additions and 51 deletions
Download Patch File Download Diff File Expand all files Collapse all files
code/training_wrapper.py
+57 -25
View File
@@ -1,8 +1,10 @@
import json
import socket
import sys
import os
import argparse
import logging
import time
from datetime import datetime
import lmdb
@@ -19,15 +21,12 @@ from utils.data_utils import LMDBIterableDataset
from utils.utils import get_variable_from_module, get_logger, convert_for_json
from utils.training import train_model
print(f"PID: {os.getpid()} on host: {socket.gethostname()}")
dotenv.load_dotenv()
logger = None
# set up distributed training
dist.init_process_group(backend="nccl", init_method="env://")
local_rank = torch.distributed.get_rank()
torch.cuda.set_device(local_rank)
def prepare_run(results_dir: str,
lmdb_root_dir: str,
@@ -59,7 +58,8 @@ def train(
train_ids: list,
val_ids: list,
dataset_dir: str,
log_dir: str) -> None:
log_dir: str,
num_dataloader_workers: int = 1) -> None:
# create training and validation loaders
logger.info("Creating training loaders")
train_loader = LMDBIterableDataset(dataset_dir,
@@ -88,14 +88,14 @@ def train(
train_dataset=train_loader,
val_dataset=val_loader,
log_dir=log_dir,
logger=logger
logger=logger,
num_dataloader_workers=num_dataloader_workers,
)
def evaluate(model_configuration: dict,
training_configuration: dict,
test_ids: list,
device) -> dict:
test_ids: list) -> dict:
logger.info(f"Evaluating model {model_configuration['id']}")
eval_functions = get_eval_functions(model_configuration)
results = evaluate_model(model_configuration,
@@ -114,6 +114,18 @@ def evaluate(model_configuration: dict,
if __name__ == "__main__":
# set up distributed training
dist.init_process_group(backend="nccl", init_method="env://")
# sleeping for a bit to allow all processes to initialize
time.sleep(5)
local_rank = torch.distributed.get_rank()
world_size = torch.distributed.get_world_size()
torch.cuda.set_device(local_rank)
print(f"Rank {local_rank} initialized with world size {world_size}")
parser = argparse.ArgumentParser(description="Training wrapper for model training")
parser.add_argument("run_configuration_module",
type=str,
@@ -144,11 +156,11 @@ if __name__ == "__main__":
required=False,
default=None,
help="Limit the number of items to process, default is None (no limit)")
parser.add_argument("--device",
type=str,
parser.add_argument("--num_dataloader_workers",
type=int,
required=False,
default="cuda",
help="Device to use for training, default is cuda")
default=1,
help="Number of workers for the dataloader, default is 1")
args = parser.parse_args()
# load run configuration from module
@@ -201,8 +213,11 @@ if __name__ == "__main__":
if not os.path.exists(log_dir):
os.makedirs(log_dir)
item_limit = args.item_limit if args.item_limit else run_configuration["item_limit"]
if item_limit == -1:
try:
item_limit = args.item_limit if args.item_limit else run_configuration["item_limit"]
if item_limit == -1:
item_limit = None
except:
item_limit = None
run_name = run_configuration["name"]
@@ -214,29 +229,42 @@ if __name__ == "__main__":
else:
run_id = None
if torch.distributed.is_initialized():
torch.distributed.barrier()
# broadcast run_id to all processes
if dist.is_initialized():
print(f"Rank {local_rank}: Broadcasting run_id {run_id}")
run_id_list = [run_id]
torch.distributed.broadcast_object_list(run_id_list, src=0)
run_id = run_id_list[0]
# make sure run_id is a string
run_id = str(run_id)
print(f"Rank {local_rank}: fetched run_id {run_id}")
# append run_id to results_dir and log_dir
results_dir = os.path.join(results_dir, run_id)
log_dir = os.path.join(log_dir, run_id)
print(f"Rank {local_rank}: Results directory: {results_dir}")
print(f"Rank {local_rank}: Log directory: {log_dir}")
# create directories if they do not exist, only on the main process
if torch.distributed.get_rank() == 0 or not dist.is_initialized():
print(f"Rank {local_rank}: Creating directories for run {run_name}")
if not os.path.exists(results_dir):
os.makedirs(results_dir)
if not os.path.exists(log_dir):
os.makedirs(log_dir)
# set up logger
logger = get_logger(module_name=run_name, filename=os.path.join(log_dir, "main.log"))
# sync all processes to make sure the directories are created
if dist.is_initialized():
print(f"Rank {local_rank}: Waiting for all processes to create directories")
dist.barrier()
# set up logger for all processes
logger = get_logger(module_name=run_name, filename=os.path.join(log_dir, f"main_{local_rank}.log"))
logger.info(f"Rank {local_rank}: Starting run {run_name}")
logger.info(f"Rank {local_rank}: Run ID: {run_id}")
@@ -270,14 +298,18 @@ if __name__ == "__main__":
base_training_configuration=run_training_configuration
)
logger.info(f"Rank {local_rank}: Finished preparing run {run_step_name}")
# sync after preparing run
if dist.is_initialized():
dist.barrier()
else:
# wait for rank 0 to finish preparing run
if dist.is_initialized():
dist.barrier()
# all ranks sync here
if dist.is_initialized():
if local_rank != 0:
logger.info(f"Rank {local_rank}: Waiting for rank 0 to finish preparing run {run_step_name}")
dist.barrier()
if local_rank != 0:
logger.info(
f"Rank {local_rank}: Finished waiting for rank 0 to finish preparing run {run_step_name}")
if local_rank != 0:
# after the barrier, rank 0 will have prepared the run
run_model_configuration, run_training_configuration, train_ids, val_ids, test_ids, dataset_dir = prepare_run(
results_dir=results_dir,
lmdb_root_dir=lmdb_root_dir,