Initial commit
This commit is contained in:
王新升
@@ -0,0 +1,501 @@
|
||||
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()
|
||||
Reference in New Issue
Block a user