| |
| import logging |
| import random |
| from typing import List, Optional, Tuple |
|
|
| import numpy as np |
| import torch |
| from torch.utils.data import DataLoader |
|
|
| from mmengine.device import is_cuda_available, is_musa_available |
| from mmengine.dist import get_rank, sync_random_seed |
| from mmengine.logging import print_log |
| from mmengine.utils import digit_version, is_list_of |
| from mmengine.utils.dl_utils import TORCH_VERSION |
|
|
|
|
| def calc_dynamic_intervals( |
| start_interval: int, |
| dynamic_interval_list: Optional[List[Tuple[int, int]]] = None |
| ) -> Tuple[List[int], List[int]]: |
| """Calculate dynamic intervals. |
| |
| Args: |
| start_interval (int): The interval used in the beginning. |
| dynamic_interval_list (List[Tuple[int, int]], optional): The |
| first element in the tuple is a milestone and the second |
| element is a interval. The interval is used after the |
| corresponding milestone. Defaults to None. |
| |
| Returns: |
| Tuple[List[int], List[int]]: a list of milestone and its corresponding |
| intervals. |
| """ |
| if dynamic_interval_list is None: |
| return [0], [start_interval] |
|
|
| assert is_list_of(dynamic_interval_list, tuple) |
|
|
| dynamic_milestones = [0] |
| dynamic_milestones.extend( |
| [dynamic_interval[0] for dynamic_interval in dynamic_interval_list]) |
| dynamic_intervals = [start_interval] |
| dynamic_intervals.extend( |
| [dynamic_interval[1] for dynamic_interval in dynamic_interval_list]) |
| return dynamic_milestones, dynamic_intervals |
|
|
|
|
| def set_random_seed(seed: Optional[int] = None, |
| deterministic: bool = False, |
| diff_rank_seed: bool = False) -> int: |
| """Set random seed. |
| |
| Args: |
| seed (int, optional): Seed to be used. |
| deterministic (bool): Whether to set the deterministic option for |
| CUDNN backend, i.e., set `torch.backends.cudnn.deterministic` |
| to True and `torch.backends.cudnn.benchmark` to False. |
| Defaults to False. |
| diff_rank_seed (bool): Whether to add rank number to the random seed to |
| have different random seed in different threads. Defaults to False. |
| """ |
| if seed is None: |
| seed = sync_random_seed() |
|
|
| if diff_rank_seed: |
| rank = get_rank() |
| seed += rank |
|
|
| random.seed(seed) |
| np.random.seed(seed) |
| torch.manual_seed(seed) |
| |
| if is_cuda_available(): |
| torch.cuda.manual_seed_all(seed) |
| elif is_musa_available(): |
| torch.musa.manual_seed_all(seed) |
| |
| if deterministic: |
| if torch.backends.cudnn.benchmark: |
| print_log( |
| 'torch.backends.cudnn.benchmark is going to be set as ' |
| '`False` to cause cuDNN to deterministically select an ' |
| 'algorithm', |
| logger='current', |
| level=logging.WARNING) |
| torch.backends.cudnn.deterministic = True |
| torch.backends.cudnn.benchmark = False |
|
|
| if digit_version(TORCH_VERSION) >= digit_version('1.10.0'): |
| torch.use_deterministic_algorithms(True) |
| return seed |
|
|
|
|
| def _get_batch_size(dataloader: dict): |
| if isinstance(dataloader, dict): |
| if 'batch_size' in dataloader: |
| return dataloader['batch_size'] |
| elif ('batch_sampler' in dataloader |
| and 'batch_size' in dataloader['batch_sampler']): |
| return dataloader['batch_sampler']['batch_size'] |
| else: |
| raise ValueError('Please set batch_size in `Dataloader` or ' |
| '`batch_sampler`') |
| elif isinstance(dataloader, DataLoader): |
| return dataloader.batch_sampler.batch_size |
| else: |
| raise ValueError('dataloader should be a dict or a Dataloader ' |
| f'instance, but got {type(dataloader)}') |
|
|
