502 lines
22 KiB
Python
502 lines
22 KiB
Python
|
import os
|
|||
|
|
import random
|
|||
|
|
import time
|
|||
|
|
import yaml
|
|||
|
|
import wandb
|
|||
|
|
import numpy as np
|
|||
|
|
import torch
|
|||
|
|
import argparse
|
|||
|
|
from typing import Dict, List, Tuple, Union
|
|||
|
|
from omegaconf import OmegaConf
|
|||
|
|
from ml_collections import ConfigDict
|
|||
|
|
import torch.distributed as dist
|
|||
|
|
from torch import nn
|
|||
|
|
|
|||
|
|
|
|||
|
|
def parse_args_train(dict_args: Union[Dict, None]) -> argparse.Namespace:
|
|||
|
|
"""
|
|||
|
|
Parse command-line arguments for training configuration.
|
|||
|
|
|
|||
|
|
This function constructs an argument parser for model, dataset, training, and logging
|
|||
|
|
options, merges overrides from a provided dictionary (if any), and returns the parsed
|
|||
|
|
arguments. If `dict_args` is None, the arguments are parsed from `sys.argv`.
|
|||
|
|
|
|||
|
|
Args:
|
|||
|
|
dict_args (Dict | None): Optional dictionary of argument overrides. Keys should
|
|||
|
|
match the defined CLI options.
|
|||
|
|
|
|||
|
|
Returns:
|
|||
|
|
argparse.Namespace: Parsed arguments namespace containing all configuration
|
|||
|
|
values required for training.
|
|||
|
|
"""
|
|||
|
|
parser = argparse.ArgumentParser()
|
|||
|
|
parser.add_argument("--model_type", type=str, default='mdx23c',
|
|||
|
|
help="One of mdx23c, htdemucs, segm_models, mel_band_roformer, bs_roformer, swin_upernet, bandit")
|
|||
|
|
parser.add_argument("--config_path", type=str, help="path to config file")
|
|||
|
|
parser.add_argument("--start_check_point", type=str, default='', help="Initial checkpoint to start training")
|
|||
|
|
parser.add_argument("--load_optimizer", action='store_true', help="Load optimizer state from checkpoint (if available)")
|
|||
|
|
parser.add_argument("--load_scheduler", action='store_true', help="Load scheduler state from checkpoint (if available)")
|
|||
|
|
parser.add_argument("--load_epoch", action='store_true', help="Load epoch number from checkpoint (if available)")
|
|||
|
|
parser.add_argument("--load_best_metric", action='store_true', help="Load best metric from checkpoint (if available)")
|
|||
|
|
parser.add_argument("--load_all_metrics", action='store_true', help="Load all metrics from checkpoint (if available)")
|
|||
|
|
parser.add_argument("--results_path", type=str,
|
|||
|
|
help="path to folder where results will be stored (weights, metadata)")
|
|||
|
|
parser.add_argument("--data_path", nargs="+", type=str, help="Dataset data paths. You can provide several folders.")
|
|||
|
|
parser.add_argument("--dataset_type", type=int, default=1,
|
|||
|
|
help="Dataset type. Must be one of: 1, 2, 3 or 4. Details here: https://github.com/ZFTurbo/Music-Source-Separation-Training/blob/main/docs/dataset_types.md")
|
|||
|
|
parser.add_argument("--valid_path", nargs="+", type=str,
|
|||
|
|
help="validation data paths. You can provide several folders.")
|
|||
|
|
parser.add_argument("--num_workers", type=int, default=0, help="dataloader num_workers")
|
|||
|
|
parser.add_argument("--pin_memory", action='store_true', help="dataloader pin_memory")
|
|||
|
|
parser.add_argument("--seed", type=int, default=0, help="random seed")
|
|||
|
|
parser.add_argument("--device_ids", nargs='+', type=int, default=[0], help='list of gpu ids')
|
|||
|
|
parser.add_argument("--loss", type=str, nargs='+', choices=['masked_loss', 'mse_loss', 'l1_loss',
|
|||
|
|
'multistft_loss', 'spec_masked_loss', 'spec_rmse_loss', 'log_wmse_loss'],
|
|||
|
|
default=['masked_loss'], help="List of loss functions to use")
|
|||
|
|
parser.add_argument("--masked_loss_coef", type=float, default=1., help="Coef for loss")
|
|||
|
|
parser.add_argument("--mse_loss_coef", type=float, default=1., help="Coef for loss")
|
|||
|
|
parser.add_argument("--l1_loss_coef", type=float, default=1., help="Coef for loss")
|
|||
|
|
parser.add_argument("--log_wmse_loss_coef", type=float, default=1., help="Coef for loss")
|
|||
|
|
parser.add_argument("--multistft_loss_coef", type=float, default=0.001, help="Coef for loss")
|
|||
|
|
parser.add_argument("--spec_masked_loss_coef", type=float, default=1, help="Coef for loss")
|
|||
|
|
parser.add_argument("--spec_rmse_loss_coef", type=float, default=1, help="Coef for loss")
|
|||
|
|
parser.add_argument("--wandb_key", type=str, default='', help='wandb API Key')
|
|||
|
|
parser.add_argument("--wandb_offline", action='store_true', help='local wandb')
|
|||
|
|
parser.add_argument("--pre_valid", action='store_true', help='Run validation before training')
|
|||
|
|
parser.add_argument("--metrics", nargs='+', type=str, default=["sdr"],
|
|||
|
|
choices=['sdr', 'l1_freq', 'si_sdr', 'log_wmse', 'aura_stft', 'aura_mrstft', 'bleedless',
|
|||
|
|
'fullness'], help='List of metrics to use.')
|
|||
|
|
parser.add_argument("--metric_for_scheduler", default="sdr",
|
|||
|
|
choices=['sdr', 'l1_freq', 'si_sdr', 'log_wmse', 'aura_stft', 'aura_mrstft', 'bleedless',
|
|||
|
|
'fullness'], help='Metric which will be used for scheduler.')
|
|||
|
|
parser.add_argument("--train_lora", action='store_true', help="Train with LoRA")
|
|||
|
|
parser.add_argument("--lora_checkpoint", type=str, default='', help="Initial checkpoint to LoRA weights")
|
|||
|
|
parser.add_argument("--each_metrics_in_name", action='store_true',
|
|||
|
|
help="All stems in naming checkpoints")
|
|||
|
|
parser.add_argument("--use_standard_loss", action='store_true',
|
|||
|
|
help="Roformers will use provided loss instead of internal")
|
|||
|
|
parser.add_argument("--save_weights_every_epoch", action='store_true',
|
|||
|
|
help="Weights will be saved every epoch with all metric values")
|
|||
|
|
parser.add_argument("--persistent_workers", action='store_true',
|
|||
|
|
help="dataloader persistent_workers")
|
|||
|
|
parser.add_argument("--prefetch_factor", type=int, default=None,
|
|||
|
|
help="dataloader prefetch_factor")
|
|||
|
|
parser.add_argument("--set_per_process_memory_fraction", action='store_true',
|
|||
|
|
help="using only VRAM, no RAM")
|
|||
|
|
|
|||
|
|
if dict_args is not None:
|
|||
|
|
args = parser.parse_args([])
|
|||
|
|
args_dict = vars(args)
|
|||
|
|
args_dict.update(dict_args)
|
|||
|
|
args = argparse.Namespace(**args_dict)
|
|||
|
|
else:
|
|||
|
|
args = parser.parse_args()
|
|||
|
|
|
|||
|
|
if args.metric_for_scheduler not in args.metrics:
|
|||
|
|
args.metrics += [args.metric_for_scheduler]
|
|||
|
|
|
|||
|
|
get_internal_loss = (args.model_type in ('mel_band_conformer',) or 'roformer' in args.model_type
|
|||
|
|
) and not args.use_standard_loss
|
|||
|
|
if get_internal_loss:
|
|||
|
|
args.loss = [f'{args.model_type}_loss']
|
|||
|
|
return args
|
|||
|
|
|
|||
|
|
|
|||
|
|
def parse_args_valid(dict_args: Union[Dict, None]) -> argparse.Namespace:
|
|||
|
|
"""
|
|||
|
|
Parse command-line arguments for validation configuration.
|
|||
|
|
|
|||
|
|
Builds the CLI for model selection, configuration paths, validation data
|
|||
|
|
locations, output/spectrogram saving options, device/runtime settings, and
|
|||
|
|
evaluation metrics. If `dict_args` is provided, its key–value pairs override
|
|||
|
|
or set the parsed arguments; otherwise arguments are read from `sys.argv`.
|
|||
|
|
|
|||
|
|
Args:
|
|||
|
|
dict_args (Union[Dict, None]): Optional mapping of argument names to values
|
|||
|
|
used to override or supply CLI options programmatically.
|
|||
|
|
|
|||
|
|
Returns:
|
|||
|
|
argparse.Namespace: Parsed arguments namespace containing all validation
|
|||
|
|
configuration values.
|
|||
|
|
"""
|
|||
|
|
parser = argparse.ArgumentParser()
|
|||
|
|
parser.add_argument("--model_type", type=str, default='mdx23c',
|
|||
|
|
help="One of mdx23c, htdemucs, segm_models, mel_band_roformer,"
|
|||
|
|
" bs_roformer, swin_upernet, bandit")
|
|||
|
|
parser.add_argument("--config_path", type=str, help="Path to config file")
|
|||
|
|
parser.add_argument("--start_check_point", type=str, default='', help="Initial checkpoint"
|
|||
|
|
" to valid weights")
|
|||
|
|
parser.add_argument("--valid_path", nargs="+", type=str, help="Validate path")
|
|||
|
|
parser.add_argument("--store_dir", type=str, default="", help="Path to store results as wav file")
|
|||
|
|
parser.add_argument("--draw_spectro", type=float, default=0,
|
|||
|
|
help="If --store_dir is set then code will generate spectrograms for resulted stems as well."
|
|||
|
|
" Value defines for how many seconds os track spectrogram will be generated.")
|
|||
|
|
parser.add_argument("--device_ids", nargs='+', type=int, default=[0], help='List of gpu ids')
|
|||
|
|
parser.add_argument("--num_workers", type=int, default=0, help="Dataloader num_workers")
|
|||
|
|
parser.add_argument("--pin_memory", action='store_true', help="Dataloader pin_memory")
|
|||
|
|
parser.add_argument("--extension", type=str, default='wav', help="Choose extension for validation")
|
|||
|
|
parser.add_argument("--use_tta", action='store_true',
|
|||
|
|
help="Flag adds test time augmentation during inference (polarity and channel inverse)."
|
|||
|
|
"While this triples the runtime, it reduces noise and slightly improves prediction quality.")
|
|||
|
|
parser.add_argument("--metrics", nargs='+', type=str, default=["sdr"],
|
|||
|
|
choices=['sdr', 'l1_freq', 'si_sdr', 'neg_log_wmse', 'aura_stft', 'aura_mrstft', 'bleedless',
|
|||
|
|
'fullness'], help='List of metrics to use.')
|
|||
|
|
parser.add_argument("--lora_checkpoint", type=str, default='', help="Initial checkpoint to LoRA weights")
|
|||
|
|
|
|||
|
|
if dict_args is not None:
|
|||
|
|
args = parser.parse_args([])
|
|||
|
|
args_dict = vars(args)
|
|||
|
|
args_dict.update(dict_args)
|
|||
|
|
args = argparse.Namespace(**args_dict)
|
|||
|
|
else:
|
|||
|
|
args = parser.parse_args()
|
|||
|
|
|
|||
|
|
return args
|
|||
|
|
|
|||
|
|
|
|||
|
|
def parse_args_inference(dict_args: Union[Dict, None]) -> argparse.Namespace:
|
|||
|
|
"""
|
|||
|
|
Parse command-line arguments for inference configuration.
|
|||
|
|
|
|||
|
|
Builds the CLI for model selection, configuration path, input/output handling,
|
|||
|
|
device/runtime options, test-time augmentation, and optional LoRA checkpoints.
|
|||
|
|
If `dict_args` is provided, its key–value pairs override or supply CLI options
|
|||
|
|
programmatically; otherwise, arguments are read from `sys.argv`.
|
|||
|
|
|
|||
|
|
Args:
|
|||
|
|
dict_args (Union[Dict, None]): Optional mapping of argument names to values
|
|||
|
|
used to override or supply CLI options programmatically.
|
|||
|
|
|
|||
|
|
Returns:
|
|||
|
|
argparse.Namespace: Parsed arguments namespace containing all inference
|
|||
|
|
configuration values.
|
|||
|
|
"""
|
|||
|
|
parser = argparse.ArgumentParser()
|
|||
|
|
parser.add_argument("--model_type", type=str, default='mdx23c',
|
|||
|
|
help="One of bandit, bandit_v2, bs_roformer, htdemucs, mdx23c, mel_band_roformer,"
|
|||
|
|
" scnet, scnet_unofficial, segm_models, swin_upernet, torchseg")
|
|||
|
|
parser.add_argument("--config_path", type=str, help="path to config file")
|
|||
|
|
parser.add_argument("--start_check_point", type=str, default='', help="Initial checkpoint to valid weights")
|
|||
|
|
parser.add_argument("--input_folder", type=str, help="folder with mixtures to process")
|
|||
|
|
parser.add_argument("--store_dir", type=str, default="", help="path to store results as wav file")
|
|||
|
|
parser.add_argument("--draw_spectro", type=float, default=0,
|
|||
|
|
help="Code will generate spectrograms for resulted stems."
|
|||
|
|
" Value defines for how many seconds os track spectrogram will be generated.")
|
|||
|
|
parser.add_argument("--device_ids", nargs='+', type=int, default=0, help='list of gpu ids')
|
|||
|
|
parser.add_argument("--extract_instrumental", action='store_true',
|
|||
|
|
help="invert vocals to get instrumental if provided")
|
|||
|
|
parser.add_argument("--disable_detailed_pbar", action='store_true', help="disable detailed progress bar")
|
|||
|
|
parser.add_argument("--force_cpu", action='store_true', help="Force the use of CPU even if CUDA is available")
|
|||
|
|
parser.add_argument("--flac_file", action='store_true', help="Output flac file instead of wav")
|
|||
|
|
parser.add_argument("--pcm_type", type=str, choices=['PCM_16', 'PCM_24'], default='PCM_24',
|
|||
|
|
help="PCM type for FLAC files (PCM_16 or PCM_24)")
|
|||
|
|
parser.add_argument("--use_tta", action='store_true',
|
|||
|
|
help="Flag adds test time augmentation during inference (polarity and channel inverse)."
|
|||
|
|
"While this triples the runtime, it reduces noise and slightly improves prediction quality.")
|
|||
|
|
parser.add_argument("--lora_checkpoint", type=str, default='', help="Initial checkpoint to LoRA weights")
|
|||
|
|
|
|||
|
|
if dict_args is not None:
|
|||
|
|
args = parser.parse_args([])
|
|||
|
|
args_dict = vars(args)
|
|||
|
|
args_dict.update(dict_args)
|
|||
|
|
args = argparse.Namespace(**args_dict)
|
|||
|
|
else:
|
|||
|
|
args = parser.parse_args()
|
|||
|
|
|
|||
|
|
return args
|
|||
|
|
|
|||
|
|
|
|||
|
|
def load_config(model_type: str, config_path: str) -> Union[ConfigDict, OmegaConf]:
|
|||
|
|
"""
|
|||
|
|
Load a model configuration from a file.
|
|||
|
|
|
|||
|
|
Based on `model_type`, returns either an OmegaConf (e.g., for 'htdemucs')
|
|||
|
|
or a YAML-parsed ConfigDict for other models.
|
|||
|
|
|
|||
|
|
Args:
|
|||
|
|
model_type (str): Model identifier that determines the loader behavior
|
|||
|
|
(e.g., 'htdemucs', 'mdx23c', etc.).
|
|||
|
|
config_path (str): Path to the configuration file (YAML/OmegaConf).
|
|||
|
|
|
|||
|
|
Returns:
|
|||
|
|
Union[ConfigDict, OmegaConf]: Loaded configuration object.
|
|||
|
|
|
|||
|
|
Raises:
|
|||
|
|
FileNotFoundError: If `config_path` does not point to an existing file.
|
|||
|
|
ValueError: If the configuration cannot be parsed or is otherwise invalid.
|
|||
|
|
"""
|
|||
|
|
try:
|
|||
|
|
with open(config_path, 'r') as f:
|
|||
|
|
if model_type == 'htdemucs':
|
|||
|
|
config = OmegaConf.load(config_path)
|
|||
|
|
else:
|
|||
|
|
config = ConfigDict(yaml.load(f, Loader=yaml.FullLoader))
|
|||
|
|
return config
|
|||
|
|
except FileNotFoundError:
|
|||
|
|
raise FileNotFoundError(f"Configuration file not found at {config_path}")
|
|||
|
|
except Exception as e:
|
|||
|
|
raise ValueError(f"Error loading configuration: {e}")
|
|||
|
|
|
|||
|
|
|
|||
|
|
def get_model_from_config(model_type: str, config_path: str) -> Tuple[nn.Module, Union[ConfigDict, OmegaConf]]:
|
|||
|
|
"""
|
|||
|
|
Load and instantiate a model using a configuration file.
|
|||
|
|
|
|||
|
|
Given a `model_type` and a path to a configuration, this function loads the
|
|||
|
|
configuration (YAML or OmegaConf) and constructs the corresponding model.
|
|||
|
|
|
|||
|
|
Args:
|
|||
|
|
model_type (str): Identifier of the model family (e.g., 'mdx23c', 'htdemucs',
|
|||
|
|
'scnet', 'mel_band_conformer', etc.).
|
|||
|
|
config_path (str): Filesystem path to the configuration file used to
|
|||
|
|
initialize the model.
|
|||
|
|
|
|||
|
|
Returns:
|
|||
|
|
Tuple[nn.Module, Union[ConfigDict, OmegaConf]]: A tuple containing the
|
|||
|
|
initialized PyTorch model and the loaded configuration object.
|
|||
|
|
|
|||
|
|
Raises:
|
|||
|
|
ValueError: If `model_type` is unknown or model initialization fails.
|
|||
|
|
FileNotFoundError: If `config_path` does not exist (may be raised by the
|
|||
|
|
underlying config loader).
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
config = load_config(model_type, config_path)
|
|||
|
|
|
|||
|
|
if model_type == 'mel_band_roformer':
|
|||
|
|
from ..modules.bs_roformer import MelBandRoformer
|
|||
|
|
model = MelBandRoformer(**dict(config.model))
|
|||
|
|
else:
|
|||
|
|
raise ValueError(f"Unknown model type: {model_type}")
|
|||
|
|
|
|||
|
|
return model, config
|
|||
|
|
|
|||
|
|
|
|||
|
|
def logging(logs: List[str], text: str, verbose_logging: bool = False) -> None:
|
|||
|
|
"""
|
|||
|
|
Print a log message and optionally append it to an in-memory list.
|
|||
|
|
|
|||
|
|
In Distributed Data Parallel (DDP) contexts, the message is printed only on
|
|||
|
|
rank 0; when DDP is uninitialized, it prints unconditionally. If
|
|||
|
|
`verbose_logging` is True, the message is also appended to `logs`.
|
|||
|
|
|
|||
|
|
Args:
|
|||
|
|
logs (List[str]): Mutable list to which the message is appended when
|
|||
|
|
`verbose_logging` is True.
|
|||
|
|
text (str): The log message to print (rank 0 only under DDP) and
|
|||
|
|
optionally store.
|
|||
|
|
verbose_logging (bool, optional): If True, append `text` to `logs`.
|
|||
|
|
Defaults to False.
|
|||
|
|
|
|||
|
|
Returns:
|
|||
|
|
None: The function prints and may mutate `logs` in place.
|
|||
|
|
"""
|
|||
|
|
if not dist.is_initialized() or dist.get_rank()==0:
|
|||
|
|
print(text)
|
|||
|
|
if verbose_logging:
|
|||
|
|
logs.append(text)
|
|||
|
|
|
|||
|
|
|
|||
|
|
def write_results_in_file(store_dir: str, logs: List[str]) -> None:
|
|||
|
|
"""
|
|||
|
|
Write accumulated log messages to a results file.
|
|||
|
|
|
|||
|
|
Creates (or overwrites) a `results.txt` file inside `store_dir` and writes
|
|||
|
|
each entry from `logs` as a separate line. In Distributed Data Parallel (DDP)
|
|||
|
|
scenarios, writing is intended to occur only on rank 0.
|
|||
|
|
|
|||
|
|
Args:
|
|||
|
|
store_dir (str): Directory path where `results.txt` will be saved.
|
|||
|
|
logs (List[str]): Ordered collection of log lines to write.
|
|||
|
|
|
|||
|
|
Returns:
|
|||
|
|
None
|
|||
|
|
"""
|
|||
|
|
if not dist.is_initialized() or dist.get_rank() == 0:
|
|||
|
|
with open(f'{store_dir}/results.txt', 'w') as out:
|
|||
|
|
for item in logs:
|
|||
|
|
out.write(item + "\n")
|
|||
|
|
|
|||
|
|
|
|||
|
|
def manual_seed(seed: int) -> None:
|
|||
|
|
"""
|
|||
|
|
Initialize random seeds for reproducibility.
|
|||
|
|
|
|||
|
|
Sets the seed across Python's `random`, NumPy, and PyTorch (CPU and CUDA)
|
|||
|
|
libraries, and updates the `PYTHONHASHSEED` environment variable. This helps
|
|||
|
|
ensure deterministic behavior where possible, though some GPU operations
|
|||
|
|
may still introduce nondeterminism.
|
|||
|
|
|
|||
|
|
Args:
|
|||
|
|
seed (int): The seed value to use for all random number generators.
|
|||
|
|
|
|||
|
|
Returns:
|
|||
|
|
None
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
random.seed(seed)
|
|||
|
|
np.random.seed(seed)
|
|||
|
|
torch.manual_seed(seed)
|
|||
|
|
torch.cuda.manual_seed(seed)
|
|||
|
|
torch.cuda.manual_seed_all(seed) # if multi-GPU
|
|||
|
|
torch.backends.cudnn.deterministic = False
|
|||
|
|
os.environ["PYTHONHASHSEED"] = str(seed)
|
|||
|
|
|
|||
|
|
|
|||
|
|
def initialize_environment(seed: int, results_path: str) -> None:
|
|||
|
|
"""
|
|||
|
|
Initialize runtime environment settings.
|
|||
|
|
|
|||
|
|
Sets random seeds for reproducibility, adjusts PyTorch cuDNN behavior,
|
|||
|
|
configures multiprocessing with the 'spawn' start method, and ensures
|
|||
|
|
the results directory exists.
|
|||
|
|
|
|||
|
|
Args:
|
|||
|
|
seed (int): Random seed value for deterministic initialization.
|
|||
|
|
results_path (str): Filesystem path to create for saving results.
|
|||
|
|
|
|||
|
|
Returns:
|
|||
|
|
None
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
manual_seed(seed)
|
|||
|
|
torch.backends.cudnn.deterministic = False
|
|||
|
|
try:
|
|||
|
|
torch.multiprocessing.set_start_method('spawn')
|
|||
|
|
except Exception as e:
|
|||
|
|
pass
|
|||
|
|
os.makedirs(results_path, exist_ok=True)
|
|||
|
|
|
|||
|
|
|
|||
|
|
def initialize_environment_ddp(rank: int, world_size: int, seed: int = 0, resuls_path: str = None) -> None:
|
|||
|
|
"""
|
|||
|
|
Initialize environment for Distributed Data Parallel (DDP) training/validation.
|
|||
|
|
|
|||
|
|
Sets up the DDP process group, seeds random number generators, configures
|
|||
|
|
multiprocessing to use the 'spawn' method, and creates a results directory
|
|||
|
|
if provided.
|
|||
|
|
|
|||
|
|
Args:
|
|||
|
|
rank (int): Rank of the current process within the DDP group.
|
|||
|
|
world_size (int): Total number of processes participating in DDP.
|
|||
|
|
seed (int, optional): Random seed for reproducibility. Defaults to 0.
|
|||
|
|
resuls_path (str, optional): Directory path to create for storing results.
|
|||
|
|
If None, no directory is created. Defaults to None.
|
|||
|
|
|
|||
|
|
Returns:
|
|||
|
|
None
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
setup_ddp(rank, world_size)
|
|||
|
|
manual_seed(seed)
|
|||
|
|
|
|||
|
|
try:
|
|||
|
|
torch.multiprocessing.set_start_method('spawn', force=True) # force=True prevent errors
|
|||
|
|
except RuntimeError as e:
|
|||
|
|
if "context has already been set" not in str(e):
|
|||
|
|
raise e
|
|||
|
|
if not(resuls_path is None):
|
|||
|
|
os.makedirs(resuls_path, exist_ok=True)
|
|||
|
|
|
|||
|
|
|
|||
|
|
def gen_wandb_name(args, config) -> str:
|
|||
|
|
"""
|
|||
|
|
Generate a descriptive name for a Weights & Biases (wandb) run.
|
|||
|
|
|
|||
|
|
Combines the model type, a dash-joined list of training instruments,
|
|||
|
|
and the current date into a single string identifier.
|
|||
|
|
|
|||
|
|
Args:
|
|||
|
|
args: Parsed arguments namespace containing at least `model_type`.
|
|||
|
|
config: Configuration object/dict with a `training.instruments` field.
|
|||
|
|
|
|||
|
|
Returns:
|
|||
|
|
str: Formatted run name in the form
|
|||
|
|
"<model_type>_[<instrument1>-<instrument2>-...]_<YYYY-MM-DD>".
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
instrum = '-'.join(config['training']['instruments'])
|
|||
|
|
time_str = time.strftime("%Y-%m-%d")
|
|||
|
|
name = '{}_[{}]_{}'.format(args.model_type, instrum, time_str)
|
|||
|
|
return name
|
|||
|
|
|
|||
|
|
|
|||
|
|
def wandb_init(args: argparse.Namespace, config: Dict, batch_size: int) -> None:
|
|||
|
|
"""
|
|||
|
|
Initialize Weights & Biases (wandb) for experiment tracking.
|
|||
|
|
|
|||
|
|
Depending on the provided arguments, sets up wandb in one of three modes:
|
|||
|
|
- Offline mode when `args.wandb_offline` is True.
|
|||
|
|
- Disabled mode when no valid `wandb_key` is provided.
|
|||
|
|
- Online mode with authentication using `args.wandb_key`.
|
|||
|
|
|
|||
|
|
Args:
|
|||
|
|
args (argparse.Namespace): Parsed arguments containing wandb options
|
|||
|
|
(`wandb_offline`, `wandb_key`, `device_ids`).
|
|||
|
|
config (Dict): Experiment configuration dictionary to log.
|
|||
|
|
batch_size (int): Training batch size to include in the run configuration.
|
|||
|
|
|
|||
|
|
Returns:
|
|||
|
|
None
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
if not dist.is_initialized() or dist.get_rank() == 0:
|
|||
|
|
if args.wandb_offline:
|
|||
|
|
wandb.init(mode='offline',
|
|||
|
|
project='msst',
|
|||
|
|
name=gen_wandb_name(args, config),
|
|||
|
|
config={'config': config, 'args': args, 'device_ids': args.device_ids, 'batch_size': batch_size}
|
|||
|
|
)
|
|||
|
|
elif args.wandb_key is None or args.wandb_key.strip() == '':
|
|||
|
|
wandb.init(mode='disabled')
|
|||
|
|
else:
|
|||
|
|
wandb.login(key=args.wandb_key)
|
|||
|
|
wandb.init(
|
|||
|
|
project='msst',
|
|||
|
|
name=gen_wandb_name(args, config),
|
|||
|
|
config={'config': config, 'args': args, 'device_ids': args.device_ids, 'batch_size': batch_size}
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
|
|||
|
|
def setup_ddp(rank: int, world_size: int) -> None:
|
|||
|
|
"""
|
|||
|
|
Initialize a Distributed Data Parallel (DDP) process group.
|
|||
|
|
|
|||
|
|
Configures environment variables for the DDP master node, attempts to
|
|||
|
|
initialize the process group with the NCCL backend (preferred for GPUs),
|
|||
|
|
and falls back to the Gloo backend if NCCL is unavailable. Also sets the
|
|||
|
|
current CUDA device to match the process rank.
|
|||
|
|
|
|||
|
|
Args:
|
|||
|
|
rank (int): Rank of the current process in the DDP group.
|
|||
|
|
world_size (int): Total number of processes participating in DDP.
|
|||
|
|
|
|||
|
|
Returns:
|
|||
|
|
None
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
os.environ['MASTER_ADDR'] = 'localhost'
|
|||
|
|
os.environ['MASTER_PORT'] = '12355' # We can change and use another
|
|||
|
|
os.environ["USE_LIBUV"] = "0"
|
|||
|
|
try:
|
|||
|
|
dist.init_process_group("nccl", rank=rank, world_size=world_size)
|
|||
|
|
except:
|
|||
|
|
dist.init_process_group("gloo", rank=rank, world_size=world_size)
|
|||
|
|
if dist.get_rank()==0:
|
|||
|
|
print(f'NCCL are not available. Using "gloo" backend.')
|
|||
|
|
|
|||
|
|
torch.cuda.set_device(rank)
|
|||
|
|
|
|||
|
|
|
|||
|
|
def cleanup_ddp() -> None:
|
|||
|
|
"""
|
|||
|
|
Finalize and clean up a Distributed Data Parallel (DDP) process group.
|
|||
|
|
|
|||
|
|
Calls `torch.distributed.destroy_process_group()` to release resources
|
|||
|
|
associated with the current DDP environment.
|
|||
|
|
|
|||
|
|
Returns:
|
|||
|
|
None
|
|||
|
|
"""
|
|||
|
|
dist.destroy_process_group()
|