### Training Pipeline Initialization

Source: https://basicsr.readthedocs.io/en/latest/_modules/basicsr/train

Sets up the entire training pipeline, starting from parsing options and configuring distributed settings. It enables cuDNN benchmarking for performance and loads any necessary resume states. The function also ensures that experiment directories and loggers are initialized correctly before training begins. This is the main entry point for starting a training process.

```python
import datetime
import logging
import math
import time
import torch
from os import path as osp

from basicsr.data import build_dataloader, build_dataset
from basicsr.data.data_sampler import EnlargedSampler
from basicsr.data.prefetch_dataloader import CPUPrefetcher, CUDAPrefetcher
from basicsr.models import build_model
from basicsr.utils import (
    AvgTimer, MessageLogger, check_resume, get_env_info,
    get_root_logger, get_time_str, init_tb_logger,
    init_wandb_logger, make_exp_dirs, mkdir_and_rename, scandir
)
from basicsr.utils.options import copy_opt_file, dict2str, parse_options




[docs]
def train_pipeline(root_path):
    # parse options, set distributed setting, set random seed
    opt, args = parse_options(root_path, is_train=True)
    opt['root_path'] = root_path

    torch.backends.cudnn.benchmark = True
    # torch.backends.cudnn.deterministic = True

    # load resume states if necessary
    resume_state = load_resume_state(opt)
    # mkdir for experiments and logger
    if resume_state is None:
        make_exp_dirs(opt)
        if opt['logger'].get('use_tb_logger') and 'debug' not in opt['name'] and opt['rank'] == 0:

```

--------------------------------

### BasicSR Model Initialization and Execution Example

Source: https://basicsr.readthedocs.io/en/latest/_modules/basicsr/archs/swinir_arch

Demonstrates how to initialize and run a SwinIR model for super-resolution. It sets up model parameters, prints the model structure, calculates and prints FLOPs, and performs a forward pass with random input data, printing the output shape.

```python
if __name__ == '__main__':
    upscale = 4
    window_size = 8
    height = (1024 // upscale // window_size + 1) * window_size
    width = (720 // upscale // window_size + 1) * window_size
    model = SwinIR(
        upscale=2,
        img_size=(height, width),
        window_size=window_size,
        img_range=1.,
        depths=[6, 6, 6, 6],
        embed_dim=60,
        num_heads=[6, 6, 6, 6],
        mlp_ratio=2,
        upsampler='pixelshuffledirect')
    print(model)
    print(height, width, model.flops() / 1e9)

    x = torch.randn((1, 3, height, width))
    x = model(x)
    print(x.shape)
```

--------------------------------

### Setup and Manage Logger - Python

Source: https://basicsr.readthedocs.io/en/latest/_modules/basicsr/utils/logger

Initializes and configures a root logger, optionally adding a file handler. It handles logger duplication and sets log levels based on process rank. Dependencies include the `logging` module and `basicsr.utils.dist_util.get_dist_info`.

```python
import logging

initialized_logger = {}
def setup_logger(logger_name='basicsr', log_file=None, log_level=logging.INFO):
    """Setup a logger. 

    Args:
        logger_name (str): root logger name. Default: 'basicsr'.
        log_file (str | None): The log filename. If specified, a FileHandler
            will be added to the root logger.
        log_level (int): The root logger level. Note that only the process of
            rank 0 is affected, while other processes will set the level to
            "Error" and be silent most of the time.

    Returns:
        logging.Logger: The root logger.
    """
    logger = logging.getLogger(logger_name)
    # if the logger has been initialized, just return it
    if logger_name in initialized_logger:
        return logger

    format_str = '%(asctime)s %(levelname)s: %(message)s'
    stream_handler = logging.StreamHandler()
    stream_handler.setFormatter(logging.Formatter(format_str))
    logger.addHandler(stream_handler)
    logger.propagate = False
    rank, _ = get_dist_info() # Assuming get_dist_info is available
    if rank != 0:
        logger.setLevel('ERROR')
    elif log_file is not None:
        logger.setLevel(log_level)
        # add file handler
        file_handler = logging.FileHandler(log_file, 'w')
        file_handler.setFormatter(logging.Formatter(format_str))
        file_handler.setLevel(log_level)
        logger.addHandler(file_handler)
    initialized_logger[logger_name] = True
    return logger

```

--------------------------------

### Get Environment and Version Info - Python

Source: https://basicsr.readthedocs.io/en/latest/_modules/basicsr/utils/logger

Retrieves and formats environment information, primarily focusing on software versions. It requires `torch`, `torchvision`, and `basicsr.version.__version__`. The output is a formatted string.

```python
import torch
import torchvision

from basicsr.version import __version__
def get_env_info():
    """Get environment information.

    Currently, only log the software version.
    """
    msg = r"""
                ____                _       _____  ____
               / __ ) ____ _ _____ (_)_____/ ___/ / __ \
              / __  |/ __ `// ___// // ___/\__ \ / /_/ /
             / /_/ // /_/ /(__  )/ // /__ ___/ // _, _/
            /_____/ \__,_//____//_/ \___//____//_/ |_|
     ______                   __   __                 __      __
    / ____/____   ____   ____/ /  / /   __  __ _____ / /__   / /
   / / __ / __ \ / __ \ / __  /  / /   / / / // ___// //_/  / /
  / /_/ // /_/ // /_/ // /_/ /  / /___/ /_/ // /__ / /<    /_
  \____/ \____/ \____/ \____/  /_____/\____/ \___//_/|_|  (_)
    """
    msg += ('\nVersion Information: '
            f'\n\tBasicSR: {__version__}'
            f'\n\tPyTorch: {torch.__version__}'
            f'\n\tTorchVision: {torchvision.__version__}')
    return msg

```

--------------------------------

### Setup SRGAN Model Optimizers

Source: https://basicsr.readthedocs.io/en/latest/_modules/basicsr/models/srgan_model

Configures and initializes the optimizers for both the generator (net_g) and discriminator (net_d) networks of the SRGAN model. It extracts optimizer types and parameters from the training configuration and appends the initialized optimizers to a list for management.

```python
    def setup_optimizers(self):
        train_opt = self.opt['train']
        # optimizer g
        optim_type = train_opt['optim_g'].pop('type')
        self.optimizer_g = self.get_optimizer(optim_type, self.net_g.parameters(), **train_opt['optim_g'])
        self.optimizers.append(self.optimizer_g)
        # optimizer d
        optim_type = train_opt['optim_d'].pop('type')
        self.optimizer_d = self.get_optimizer(optim_type, self.net_d.parameters(), **train_opt['optim_d'])
        self.optimizers.append(self.optimizer_d)

```

--------------------------------

### Setup Optimizers for Generator and Discriminator

Source: https://basicsr.readthedocs.io/en/latest/_modules/basicsr/models/video_recurrent_gan_model

Configures and initializes optimizers for the generator (net_g) and discriminator (net_d). It supports fixing specific parameters, such as those in the 'spynet', by assigning different learning rates. Dependencies include get_optimizer. Inputs are training options and model parameters. Outputs are the configured optimizers.

```python
    def setup_optimizers(self):
        train_opt = self.opt['train']
        if train_opt['fix_flow']:
            normal_params = []
            flow_params = []
            for name, param in self.net_g.named_parameters():
                if 'spynet' in name:  # The fix_flow now only works for spynet.
                    flow_params.append(param)
                else:
                    normal_params.append(param)

            optim_params = [
                {  # add flow params first
                    'params': flow_params,
                    'lr': train_opt['lr_flow']
                },
                {
                    'params': normal_params,
                    'lr': train_opt['optim_g']['lr']
                },
            ]
        else:
            optim_params = self.net_g.parameters()

        # optimizer g
        optim_type = train_opt['optim_g'].pop('type')
        self.optimizer_g = self.get_optimizer(optim_type, optim_params, **train_opt['optim_g'])
        self.optimizers.append(self.optimizer_g)
        # optimizer d
        optim_type = train_opt['optim_d'].pop('type')
        self.optimizer_d = self.get_optimizer(optim_type, self.net_d.parameters(), **train_opt['optim_d'])
        self.optimizers.append(self.optimizer_d)

```

--------------------------------

### Setup Learning Rate Schedulers

Source: https://basicsr.readthedocs.io/en/latest/_modules/basicsr/models/base_model

Placeholder method for setting up learning rate schedulers. This function is intended to be overridden by subclasses to configure specific scheduler behavior based on the training options.

```python
def setup_schedulers(self):
        """Set up schedulers."""
        train_opt = self.opt['train']

```

--------------------------------

### Optimizer Setup for Generator and Discriminator

Source: https://basicsr.readthedocs.io/en/latest/_modules/basicsr/models/hifacegan_model

Sets up the optimizers for both the generator (net_g) and discriminator (net_d) networks. It extracts optimizer type and parameters from the training options and initializes the optimizers, adding them to a list of all optimizers used in the model.

```python
    def setup_optimizers(self):
        train_opt = self.opt['train']
        # optimizer g
        optim_type = train_opt['optim_g'].pop('type')
        self.optimizer_g = self.get_optimizer(optim_type, self.net_g.parameters(), **train_opt['optim_g'])
        self.optimizers.append(self.optimizer_g)
        # optimizer d
        optim_type = train_opt['optim_d'].pop('type')
        self.optimizer_d = self.get_optimizer(optim_type, self.net_d.parameters(), **train_opt['optim_d'])
        self.optimizers.append(self.optimizer_d)
```

--------------------------------

### Setup StyleGAN2 Optimizers for Generator and Discriminator

Source: https://basicsr.readthedocs.io/en/latest/_modules/basicsr/models/stylegan2_model

Configures and sets up the optimizers for the generator (net_g) and discriminator (net_d) networks. It handles different parameter groups for the generator, including special treatment for StyleGAN2GeneratorC, and defines the optimizer type and learning rate.

```python
def setup_optimizers(self):
        train_opt = self.opt['train']
        # optimizer g
        net_g_reg_ratio = self.net_g_reg_every / (self.net_g_reg_every + 1)
        if self.opt['network_g']['type'] == 'StyleGAN2GeneratorC':
            normal_params = []
            style_mlp_params = []
            modulation_conv_params = []
            for name, param in self.net_g.named_parameters():
                if 'modulation' in name:
                    normal_params.append(param)
                elif 'style_mlp' in name:
                    style_mlp_params.append(param)
                elif 'modulated_conv' in name:
                    modulation_conv_params.append(param)
                else:
                    normal_params.append(param)
            optim_params_g = [
                {  # add normal params first
                    'params': normal_params,
                    'lr': train_opt['optim_g']['lr']
                },
                {
                    'params': style_mlp_params,
                    'lr': train_opt['optim_g']['lr'] * 0.01
                },
                {
                    'params': modulation_conv_params,
                    'lr': train_opt['optim_g']['lr'] / 3
                }
            ]
        else:
            normal_params = []
            for name, param in self.net_g.named_parameters():
                normal_params.append(param)
            optim_params_g = [{  # add normal params first
                'params': normal_params,
                'lr': train_opt['optim_g']['lr']
            }]

        optim_type = train_opt['optim_g'].pop('type')

```

--------------------------------

### DFDNet Model Conversion Example

Source: https://basicsr.readthedocs.io/en/latest/_modules/scripts/model_conversion/convert_dfdnet

Demonstrates the practical application of the `convert_net` function. It loads an official original DFDNet model, initializes a new DFDNet instance, converts the weights using the `convert_net` function, and saves the converted model to a new checkpoint file.

```python
if __name__ == '__main__':
    ori_net = torch.load('experiments/pretrained_models/DFDNet/DFDNet_official_original.pth')
    dfd_net = DFDNet(64, dict_path='experiments/pretrained_models/DFDNet/DFDNet_dict_512.pth')
    crt_net = dfd_net.state_dict()
    crt_net_params = convert_net(ori_net, crt_net)

    torch.save(
        dict(params=crt_net_params),
        'experiments/pretrained_models/DFDNet/DFDNet_official.pth',
        _use_new_zipfile_serialization=False)

```

--------------------------------

### Abstract Base Storage Backend Definition

Source: https://basicsr.readthedocs.io/en/latest/_modules/basicsr/utils/file_client

Defines the abstract interface for storage backends, requiring implementations for `get()` (byte stream) and `get_text()` (text stream) methods. It serves as a blueprint for all concrete storage backend classes.

```python
from abc import ABCMeta, abstractmethod


[docs]class BaseStorageBackend(metaclass=ABCMeta):
    """Abstract class of storage backends.

    All backends need to implement two apis: ``get()`` and ``get_text()``.
    ``get()`` reads the file as a byte stream and ``get_text()`` reads the file
    as texts.
    """


    @abstractmethod
    def get(self, filepath):
        pass



    @abstractmethod
    def get_text(self, filepath):
        pass

```

--------------------------------

### EDVRModel Initialization and Optimizer Setup in Python

Source: https://basicsr.readthedocs.io/en/latest/_modules/basicsr/models/edvr_model

Initializes the EDVRModel, handling potential TSA (Temporal Stable Attention) module training iterations. It also sets up the optimizer, with special handling for deformable convolutional layer learning rates.

```python
from basicsr.utils import get_root_logger
from basicsr.utils.registry import MODEL_REGISTRY
from .video_base_model import VideoBaseModel



[docs]@MODEL_REGISTRY.register()
class EDVRModel(VideoBaseModel):
    """EDVR Model.

    Paper: EDVR: Video Restoration with Enhanced Deformable Convolutional Networks.  # noqa: E501
    """

    def __init__(self, opt):
        super(EDVRModel, self).__init__(opt)
        if self.is_train:
            self.train_tsa_iter = opt['train'].get('tsa_iter')


[docs]    def setup_optimizers(self):
        train_opt = self.opt['train']
        dcn_lr_mul = train_opt.get('dcn_lr_mul', 1)
        logger = get_root_logger()
        logger.info(f'Multiple the learning rate for dcn with {dcn_lr_mul}.')
        if dcn_lr_mul == 1:
            optim_params = self.net_g.parameters()
        else:  # separate dcn params and normal params for different lr
            normal_params = []
            dcn_params = []
            for name, param in self.net_g.named_parameters():
                if 'dcn' in name:
                    dcn_params.append(param)
                else:
                    normal_params.append(param)
            optim_params = [
                {
                    'params': normal_params,
                    'lr': train_opt['optim_g']['lr']
                },
                {
                    'params': dcn_params,
                    'lr': train_opt['optim_g']['lr'] * dcn_lr_mul
                },
            ]

        optim_type = train_opt['optim_g'].pop('type')
        self.optimizer_g = self.get_optimizer(optim_type, optim_params, **train_opt['optim_g'])
        self.optimizers.append(self.optimizer_g)



[docs]    def optimize_parameters(self, current_iter):
        if self.train_tsa_iter:
            if current_iter == 1:
                logger = get_root_logger()
                logger.info(f'Only train TSA module for {self.train_tsa_iter} iters.')
                for name, param in self.net_g.named_parameters():
                    if 'fusion' not in name:
                        param.requires_grad = False
            elif current_iter == self.train_tsa_iter:
                logger = get_root_logger()
                logger.warning('Train all the parameters.')
                for param in self.net_g.parameters():
                    param.requires_grad = True

        super(EDVRModel, self).optimize_parameters(current_iter)

```

--------------------------------

### Setup Optimizer for Generator Network (Python)

Source: https://basicsr.readthedocs.io/en/latest/_modules/basicsr/models/sr_model

Configures the optimizer for the generator network (net_g). It extracts optimization parameters from the network, filters out parameters that do not require gradients, and initializes the optimizer based on the 'optim_g' configuration provided in the training options. This function assumes 'get_optimizer' and 'setup_schedulers' are available from the base class.

```Python
    def setup_optimizers(self):
        train_opt = self.opt['train']
        optim_params = []
        for k, v in self.net_g.named_parameters():
            if v.requires_grad:
                optim_params.append(v)
            else:
                logger = get_root_logger()
                logger.warning(f'Params {k} will not be optimized.')

        optim_type = train_opt['optim_g'].pop('type')
        self.optimizer_g = self.get_optimizer(optim_type, optim_params, **train_opt['optim_g'])
        self.optimizers.append(self.optimizer_g)

```

--------------------------------

### LMDB Storage Backend Implementation

Source: https://basicsr.readthedocs.io/en/latest/_modules/basicsr/utils/file_client

Implements the `BaseStorageBackend` using LMDB (Lightning Memory-Mapped Database). It supports multiple database paths and client keys, with options for read-only access, locking, and readahead. Only the `get()` method is implemented for retrieving data by key.

```python
class LmdbBackend(BaseStorageBackend):
    """Lmdb storage backend.

    Args:
        db_paths (str | list[str]): Lmdb database paths.
        client_keys (str | list[str]): Lmdb client keys. Default: 'default'.
        readonly (bool, optional): Lmdb environment parameter. If True,
            disallow any write operations. Default: True.
        lock (bool, optional): Lmdb environment parameter. If False, when
            concurrent access occurs, do not lock the database. Default: False.
        readahead (bool, optional): Lmdb environment parameter. If False,
            disable the OS filesystem readahead mechanism, which may improve
            random read performance when a database is larger than RAM.
            Default: False.

    Attributes:
        db_paths (list): Lmdb database path.
        _client (list): A list of several lmdb envs.
    """

    def __init__(self, db_paths, client_keys='default', readonly=True, lock=False, readahead=False, **kwargs):
        try:
            import lmdb
        except ImportError:
            raise ImportError('Please install lmdb to enable LmdbBackend.')

        if isinstance(client_keys, str):
            client_keys = [client_keys]

        if isinstance(db_paths, list):
            self.db_paths = [str(v) for v in db_paths]
        elif isinstance(db_paths, str):
            self.db_paths = [str(db_paths)]
        assert len(client_keys) == len(self.db_paths), ('client_keys and db_paths should have the same length, '
                                                        f'but received {len(client_keys)} and {len(self.db_paths)}.')

        self._client = {}
        for client, path in zip(client_keys, self.db_paths):
            self._client[client] = lmdb.open(path, readonly=readonly, lock=lock, readahead=readahead, **kwargs)


    def get(self, filepath, client_key):
        """Get values according to the filepath from one lmdb named client_key.

        Args:
            filepath (str | obj:`Path`): Here, filepath is the lmdb key.
            client_key (str): Used for distinguishing different lmdb envs.
        """
        filepath = str(filepath)
        assert client_key in self._client, (f'client_key {client_key} is not in lmdb clients.')
        client = self._client[client_key]
        with client.begin(write=False) as txn:
            value_buf = txn.get(filepath.encode('ascii'))
        return value_buf



    def get_text(self, filepath):
        raise NotImplementedError

```

--------------------------------

### Setup Optimizers for VideoRecurrentModel in PyTorch

Source: https://basicsr.readthedocs.io/en/latest/_modules/basicsr/models/video_recurrent_model

Configures optimizers for the recurrent video model. It supports separate learning rates for different network components, such as the flow network, by adjusting optimizer parameters based on the 'flow_lr_mul' setting.

```python
def setup_optimizers(self):
        train_opt = self.opt['train']
        flow_lr_mul = train_opt.get('flow_lr_mul', 1)
        logger = get_root_logger()
        logger.info(f'Multiple the learning rate for flow network with {flow_lr_mul}.')
        if flow_lr_mul == 1:
            optim_params = self.net_g.parameters()
        else:
            normal_params = []
            flow_params = []
            for name, param in self.net_g.named_parameters():
                if 'spynet' in name:
                    flow_params.append(param)
                else:
                    normal_params.append(param)
            optim_params = [
                {
                    'params': normal_params,
                    'lr': train_opt['optim_g']['lr']
                },
                {
                    'params': flow_params,
                    'lr': train_opt['optim_g']['lr'] * flow_lr_mul
                },
            ]

        optim_type = train_opt['optim_g'].pop('type')
        self.optimizer_g = self.get_optimizer(optim_type, optim_params, **train_opt['optim_g'])
        self.optimizers.append(self.optimizer_g)

```

--------------------------------

### Override Autograd Function Forward and Setup Context in Python

Source: https://basicsr.readthedocs.io/en/latest/api/basicsr.ops.fused_act

Demonstrates how to override the `forward` and `setup_context` static methods for custom autograd functions in PyTorch. The `forward` method no longer accepts `ctx`, requiring `setup_context` to handle context initialization with inputs and outputs.

```python
@staticmethod
def forward(*args: Any, **kwargs: Any) -> Any:
    pass

@staticmethod
def setup_context(ctx: Any, inputs: Tuple[Any, ...], output: Any) -> None:
    pass

```

--------------------------------

### Resume PyTorch Training

Source: https://basicsr.readthedocs.io/en/latest/_modules/basicsr/models/base_model

Reloads optimizer and scheduler states from a given resume state dictionary to continue training. It asserts that the number of optimizers and schedulers in the resume state matches the current training setup.

```python
    def resume_training(self, resume_state):
        """Reload the optimizers and schedulers for resumed training.

        Args:
            resume_state (dict): Resume state.
        """
        resume_optimizers = resume_state['optimizers']
        resume_schedulers = resume_state['schedulers']
        assert len(resume_optimizers) == len(self.optimizers), 'Wrong lengths of optimizers'
        assert len(resume_schedulers) == len(self.schedulers), 'Wrong lengths of schedulers'
        for i, o in enumerate(resume_optimizers):
            self.optimizers[i].load_state_dict(o)
```

--------------------------------

### Hard Disk Storage Backend Implementation

Source: https://basicsr.readthedocs.io/en/latest/_modules/basicsr/utils/file_client

Implements the `BaseStorageBackend` for direct file system access. It provides `get()` for reading files as byte streams and `get_text()` for reading files as text strings using standard Python file I/O.

```python
class HardDiskBackend(BaseStorageBackend):
    """Raw hard disks storage backend."""


    def get(self, filepath):
        filepath = str(filepath)
        with open(filepath, 'rb') as f:
            value_buf = f.read()
        return value_buf



    def get_text(self, filepath):
        filepath = str(filepath)
        with open(filepath, 'r') as f:
            value_buf = f.read()
        return value_buf

```

--------------------------------

### Distributed Validation Setup in BasicSR

Source: https://basicsr.readthedocs.io/en/latest/_modules/basicsr/models/stylegan2_model

Handles the setup for distributed validation. If the current process is the main process (rank 0), it calls the nondistributed validation function.

```Python
def dist_validation(self, dataloader, current_iter, tb_logger, save_img):
        if self.opt['rank'] == 0:
            self.nondist_validation(dataloader, current_iter, tb_logger, save_img)
```

--------------------------------

### Build Dataset from Options - basicsr.data

Source: https://basicsr.readthedocs.io/en/latest/_modules/basicsr/data/__init__

Constructs a dataset instance based on provided configuration options. It expects a dictionary with 'type' and 'name' keys. Dependencies include the DATASET_REGISTRY for retrieving dataset classes and get_root_logger for logging. Outputs a dataset object.

```python
def build_dataset(dataset_opt):
    """Build dataset from options.

    Args:
        dataset_opt (dict): Configuration for dataset. It must contain:
            name (str): Dataset name.
            type (str): Dataset type.
    """
    dataset_opt = deepcopy(dataset_opt)
    dataset = DATASET_REGISTRY.get(dataset_opt['type'])(dataset_opt)
    logger = get_root_logger()
    logger.info(f'Dataset [{dataset.__class__.__name__}] - {dataset_opt["name"]} is built.')
    return dataset
```

--------------------------------

### RealESRNetModel Initialization and Enhancements - Python

Source: https://basicsr.readthedocs.io/en/latest/_modules/basicsr/models/realesrnet_model

Initializes the RealESRNetModel, setting up JPEG compression and USM sharpening tools. It also configures a queue size for managing training data diversity.

```python
import numpy as np
import random
import torch
from torch.nn import functional as F

from basicsr.data.degradations import random_add_gaussian_noise_pt, random_add_poisson_noise_pt
from basicsr.data.transforms import paired_random_crop
from basicsr.models.sr_model import SRModel
from basicsr.utils import DiffJPEG, USMSharp
from basicsr.utils.img_process_util import filter2D
from basicsr.utils.registry import MODEL_REGISTRY



@MODEL_REGISTRY.register(suffix='basicsr')
class RealESRNetModel(SRModel):
    """RealESRNet Model for Real-ESRGAN: Training Real-World Blind Super-Resolution with Pure Synthetic Data.

    It is trained without GAN losses.
    It mainly performs:
    1. randomly synthesize LQ images in GPU tensors
    2. optimize the networks with GAN training.
    """

    def __init__(self, opt):
        super(RealESRNetModel, self).__init__(opt)
        self.jpeger = DiffJPEG(differentiable=False).cuda()  # simulate JPEG compression artifacts
        self.usm_sharpener = USMSharp().cuda()  # do usm sharpening
        self.queue_size = opt.get('queue_size', 180)
```

--------------------------------

### Get Dataset Length

Source: https://basicsr.readthedocs.io/en/latest/_modules/basicsr/data/vimeo90k_dataset

Returns the total number of keys (e.g., image sequences) in the dataset. This method is essential for iterating through the dataset during training or evaluation.

```python
    def __len__(self):
        return len(self.keys)
```

--------------------------------

### Configure Experiment Paths for Training and Testing (Python)

Source: https://basicsr.readthedocs.io/en/latest/_modules/basicsr/utils/options

This function processes configuration options to set up directories for experiments, models, training states, logs, and visualizations. It differentiates between training and testing modes, creating appropriate root paths and subdirectories. Dependencies include the 'os.path' module.

```python
def setup_paths(opt, root_path, is_train, args=None):
    # ... (previous code for dataset and opt path expansion) ...

    if is_train:
        experiments_root = opt['path'].get('experiments_root')
        if experiments_root is None:
            experiments_root = osp.join(root_path, 'experiments')
        experiments_root = osp.join(experiments_root, opt['name'])

        opt['path']['experiments_root'] = experiments_root
        opt['path']['models'] = osp.join(experiments_root, 'models')
        opt['path']['training_states'] = osp.join(experiments_root, 'training_states')
        opt['path']['log'] = experiments_root
        opt['path']['visualization'] = osp.join(experiments_root, 'visualization')

        # change some options for debug mode
        if 'debug' in opt['name']:
            if 'val' in opt:
                opt['val']['val_freq'] = 8
            opt['logger']['print_freq'] = 1
            opt['logger']['save_checkpoint_freq'] = 8
    else:  # test
        results_root = opt['path'].get('results_root')
        if results_root is None:
            results_root = osp.join(root_path, 'results')
        results_root = osp.join(results_root, opt['name'])

        opt['path']['results_root'] = results_root
        opt['path']['log'] = results_root
        opt['path']['visualization'] = osp.join(results_root, 'visualization')

    return opt, args
```

--------------------------------

### Get Registry Keys (Python)

Source: https://basicsr.readthedocs.io/en/latest/api/basicsr.utils

Provides a method to retrieve all registered keys (names) within a registry instance. This is useful for inspecting the available components or for debugging.

```python
BACKBONE_REGISTRY.keys()
```

--------------------------------

### Get Object from Registry (Python)

Source: https://basicsr.readthedocs.io/en/latest/api/basicsr.utils

Illustrates how to retrieve a registered object from a registry by its name. This method supports retrieving objects associated with a specific suffix, defaulting to 'basicsr'.

```python
BACKBONE_REGISTRY.get(_name='MyBackbone', _suffix='basicsr')
```

--------------------------------

### Get Length of Keys

Source: https://basicsr.readthedocs.io/en/latest/_modules/basicsr/data/reds_dataset

Returns the total number of keys stored in the instance. This method is typically used to determine the size of a dataset or collection.

```python
def __len__(self):
        return len(self.keys)
```

--------------------------------

### Initialize Distributed Training Environment (Python)

Source: https://basicsr.readthedocs.io/en/latest/_modules/basicsr/utils/dist_util

Initializes the distributed training environment. Supports 'pytorch' and 'slurm' launchers. It sets up process groups for distributed communication. Dependencies include `torch`, `torch.distributed`, `torch.multiprocessing`, `os`, and `subprocess`.

```python
# Modified from https://github.com/open-mmlab/mmcv/blob/master/mmcv/runner/dist_utils.py  # noqa: E501
import functools
import os
import subprocess
import torch
import torch.distributed as dist
import torch.multiprocessing as mp



[docs]
def init_dist(launcher, backend='nccl', **kwargs):
    if mp.get_start_method(allow_none=True) is None:
        mp.set_start_method('spawn')
    if launcher == 'pytorch':
        _init_dist_pytorch(backend, **kwargs)
    elif launcher == 'slurm':
        _init_dist_slurm(backend, **kwargs)
    else:
        raise ValueError(f'Invalid launcher type: {launcher}')


def _init_dist_pytorch(backend, **kwargs):
    rank = int(os.environ['RANK'])
    num_gpus = torch.cuda.device_count()
    torch.cuda.set_device(rank % num_gpus)
    dist.init_process_group(backend=backend, **kwargs)


def _init_dist_slurm(backend, port=None):
    """Initialize slurm distributed training environment.

    If argument ``port`` is not specified, then the master port will be system
    environment variable ``MASTER_PORT``. If ``MASTER_PORT`` is not in system
    environment variable, then a default port ``29500`` will be used.

    Args:
        backend (str): Backend of torch.distributed.
        port (int, optional): Master port. Defaults to None.
    """
    proc_id = int(os.environ['SLURM_PROCID'])
    ntasks = int(os.environ['SLURM_NTASKS'])
    node_list = os.environ['SLURM_NODELIST']
    num_gpus = torch.cuda.device_count()
    torch.cuda.set_device(proc_id % num_gpus)
    addr = subprocess.getoutput(f'scontrol show hostname {node_list} | head -n1')
    # specify master port
    if port is not None:
        os.environ['MASTER_PORT'] = str(port)
    elif 'MASTER_PORT' in os.environ:
        pass  # use MASTER_PORT in the environment variable
    else:
        # 29500 is torch.distributed default port
        os.environ['MASTER_PORT'] = '29500'
    os.environ['MASTER_ADDR'] = addr
    os.environ['WORLD_SIZE'] = str(ntasks)
    os.environ['LOCAL_RANK'] = str(proc_id % num_gpus)
    os.environ['RANK'] = str(proc_id)
    dist.init_process_group(backend=backend)



[docs]
def get_dist_info():
    if dist.is_available():
        initialized = dist.is_initialized()
    else:
        initialized = False
    if initialized:
        rank = dist.get_rank()
        world_size = dist.get_world_size()
    else:
        rank = 0
        world_size = 1
    return rank, world_size




[docs]
def master_only(func):

    @functools.wraps(func)
    def wrapper(*args, **kwargs):
        rank, _ = get_dist_info()
        if rank == 0:
            return func(*args, **kwargs)

    return wrapper

```

--------------------------------

### Build Dataset API

Source: https://basicsr.readthedocs.io/en/latest/_modules/basicsr/data/__init__

Constructs a dataset instance based on the provided configuration options. It requires a 'name' and 'type' to identify and instantiate the correct dataset class from the registry.

```APIDOC
## POST /datasets

### Description
Builds a dataset instance from configuration options.

### Method
POST

### Endpoint
/datasets

### Parameters
#### Request Body
- **dataset_opt** (dict) - Required - Configuration for the dataset. Must contain 'name' (str) and 'type' (str).

### Request Example
```json
{
  "dataset_opt": {
    "name": "MyDataset",
    "type": "CustomDataset",
    "dataroot": "/path/to/data"
  }
}
```

### Response
#### Success Response (200)
- **dataset** (object) - The constructed dataset object.

#### Response Example
```json
{
  "dataset": "<dataset_object_representation>"
}
```
```

--------------------------------

### Get Current Learning Rate

Source: https://basicsr.readthedocs.io/en/latest/_modules/basicsr/models/base_model

Retrieves the current learning rate for the first optimizer's parameter groups. This is useful for monitoring the learning rate during training.

```python
def get_current_learning_rate(self):
    return [param_group['lr'] for param_group in self.optimizers[0].param_groups]
```

--------------------------------

### Initialize Logging and Environment Info

Source: https://basicsr.readthedocs.io/en/latest/_modules/basicsr/train

Sets up the root logger for the training process, including file logging and environment information. It handles the creation of log directories and specifies the log level. Dependencies include the `logging` module and utility functions like `get_time_str`, `get_env_info`, and `dict2str`.

```python
log_file = osp.join(opt['path']['log'], f"train_{opt['name']}_{get_time_str()}.log")
logger = get_root_logger(logger_name='basicsr', log_level=logging.INFO, log_file=log_file)
logger.info(get_env_info())
logger.info(dict2str(opt))
```

--------------------------------

### Get Current Time String in Python

Source: https://basicsr.readthedocs.io/en/latest/_modules/basicsr/utils/misc

Returns the current time formatted as a string in the 'YYYYMMDD_HHMMSS' format. This utility is useful for creating unique filenames or timestamps for logging and archiving.

```python
import time

def get_time_str():
    """Get current time string."""
    return time.strftime('%Y%m%d_%H%M%S', time.localtime())
```

--------------------------------

### Initialize SRModel with Network and Training Settings (Python)

Source: https://basicsr.readthedocs.io/en/latest/_modules/basicsr/models/sr_model

Initializes the SRModel, building the generator network, loading pre-trained weights if specified, and setting up training-specific components like EMA, loss functions, and optimizers. It requires the 'network_g' configuration and optionally 'path' for pre-trained models.

```Python
import torch
from collections import OrderedDict
from os import path as osp
from tqdm import tqdm

from basicsr.archs import build_network
from basicsr.losses import build_loss
from basicsr.metrics import calculate_metric
from basicsr.utils import get_root_logger, imwrite, tensor2img
from basicsr.utils.registry import MODEL_REGISTRY
from .base_model import BaseModel



@MODEL_REGISTRY.register()
class SRModel(BaseModel):
    """Base SR model for single image super-resolution."""

    def __init__(self, opt):
        super(SRModel, self).__init__(opt)

        # define network
        self.net_g = build_network(opt['network_g'])
        self.net_g = self.model_to_device(self.net_g)
        self.print_network(self.net_g)

        # load pretrained models
        load_path = self.opt['path'].get('pretrain_network_g', None)
        if load_path is not None:
            param_key = self.opt['path'].get('param_key_g', 'params')
            self.load_network(self.net_g, load_path, self.opt['path'].get('strict_load_g', True), param_key)

        if self.is_train:
            self.init_training_settings()


    def init_training_settings(self):
        self.net_g.train()
        train_opt = self.opt['train']

        self.ema_decay = train_opt.get('ema_decay', 0)
        if self.ema_decay > 0:
            logger = get_root_logger()
            logger.info(f'Use Exponential Moving Average with decay: {self.ema_decay}')
            # define network net_g with Exponential Moving Average (EMA)
            # net_g_ema is used only for testing on one GPU and saving
            # There is no need to wrap with DistributedDataParallel
            self.net_g_ema = build_network(self.opt['network_g']).to(self.device)
            # load pretrained model
            load_path = self.opt['path'].get('pretrain_network_g', None)
            if load_path is not None:
                self.load_network(self.net_g_ema, load_path, self.opt['path'].get('strict_load_g', True), 'params_ema')
            else:
                self.model_ema(0)  # copy net_g weight
            self.net_g_ema.eval()

        # define losses
        if train_opt.get('pixel_opt'):
            self.cri_pix = build_loss(train_opt['pixel_opt']).to(self.device)
        else:
            self.cri_pix = None

        if train_opt.get('perceptual_opt'):
            self.cri_perceptual = build_loss(train_opt['perceptual_opt']).to(self.device)
        else:
            self.cri_perceptual = None

        if self.cri_pix is None and self.cri_perceptual is None:
            raise ValueError('Both pixel and perceptual losses are None.')

        # set up optimizers and schedulers
        self.setup_optimizers()
        self.setup_schedulers()

```

--------------------------------

### Get Text Content from Storage Backend (Python)

Source: https://basicsr.readthedocs.io/en/latest/_modules/basicsr/utils/file_client

Retrieves the content of a file as text from the configured storage backend. This method directly calls the get_text method of the underlying backend client.

```python
    def get_text(self, filepath):
        return self.client.get_text(filepath)
```

--------------------------------

### RealESRGAN Model Initialization - Python

Source: https://basicsr.readthedocs.io/en/latest/_modules/basicsr/models/realesrgan_model

Initializes the RealESRGAN model, inheriting from SRGANModel. It sets up JPEG compression and USM sharpening utilities, and defines a queue size for managing training data diversity.

```python
import numpy as np
import random
import torch
from collections import OrderedDict
from torch.nn import functional as F

from basicsr.data.degradations import random_add_gaussian_noise_pt, random_add_poisson_noise_pt
from basicsr.data.transforms import paired_random_crop
from basicsr.losses.loss_util import get_refined_artifact_map
from basicsr.models.srgan_model import SRGANModel
from basicsr.utils import DiffJPEG, USMSharp
from basicsr.utils.img_process_util import filter2D
from basicsr.utils.registry import MODEL_REGISTRY



@MODEL_REGISTRY.register(suffix='basicsr')
class RealESRGANModel(SRGANModel):
    """RealESRGAN Model for Real-ESRGAN: Training Real-World Blind Super-Resolution with Pure Synthetic Data.

    It mainly performs:
    1. randomly synthesize LQ images in GPU tensors
    2. optimize the networks with GAN training.
    """

    def __init__(self, opt):
        super(RealESRGANModel, self).__init__(opt)
        self.jpeger = DiffJPEG(differentiable=False).cuda()  # simulate JPEG compression artifacts
        self.usm_sharpener = USMSharp().cuda()  # do usm sharpening
        self.queue_size = opt.get('queue_size', 180)

```

--------------------------------

### Memcached Storage Backend Implementation

Source: https://basicsr.readthedocs.io/en/latest/_modules/basicsr/utils/file_client

Implements the `BaseStorageBackend` using Memcached. It requires `mc` library and provides a `get()` method to retrieve data as a byte stream. The `get_text()` method is not implemented.

```python
class MemcachedBackend(BaseStorageBackend):
    """Memcached storage backend.

    Attributes:
        server_list_cfg (str): Config file for memcached server list.
        client_cfg (str): Config file for memcached client.
        sys_path (str | None): Additional path to be appended to `sys.path`.
            Default: None.
    """

    def __init__(self, server_list_cfg, client_cfg, sys_path=None):
        if sys_path is not None:
            import sys
            sys.path.append(sys_path)
        try:
            import mc
        except ImportError:
            raise ImportError('Please install memcached to enable MemcachedBackend.')

        self.server_list_cfg = server_list_cfg
        self.client_cfg = client_cfg
        self._client = mc.MemcachedClient.GetInstance(self.server_list_cfg, self.client_cfg)
        # mc.pyvector servers as a point which points to a memory cache
        self._mc_buffer = mc.pyvector()


    def get(self, filepath):
        filepath = str(filepath)
        import mc
        self._client.Get(filepath, self._mc_buffer)
        value_buf = mc.ConvertBuffer(self._mc_buffer)
        return value_buf



    def get_text(self, filepath):
        raise NotImplementedError

```

--------------------------------

### Build Dataloader - basicsr.data

Source: https://basicsr.readthedocs.io/en/latest/_modules/basicsr/data/__init__

Creates a dataloader for a given dataset and configuration. It handles different configurations for training (distributed and non-distributed) and validation/testing phases. Supports custom samplers and worker initialization for reproducibility. Dependencies include get_dist_info, PrefetchDataLoader, and torch.utils.data.DataLoader.

```python
def build_dataloader(dataset, dataset_opt, num_gpu=1, dist=False, sampler=None, seed=None):
    """Build dataloader.

    Args:
        dataset (torch.utils.data.Dataset): Dataset.
        dataset_opt (dict): Dataset options. It contains the following keys:
            phase (str): 'train' or 'val'.
            num_worker_per_gpu (int): Number of workers for each GPU.
            batch_size_per_gpu (int): Training batch size for each GPU.
        num_gpu (int): Number of GPUs. Used only in the train phase.
            Default: 1.
        dist (bool): Whether in distributed training. Used only in the train
            phase. Default: False.
        sampler (torch.utils.data.sampler): Data sampler. Default: None.
        seed (int | None): Seed. Default: None
    """
    phase = dataset_opt['phase']
    rank, _ = get_dist_info()
    if phase == 'train':
        if dist:  # distributed training
            batch_size = dataset_opt['batch_size_per_gpu']
            num_workers = dataset_opt['num_worker_per_gpu']
        else:  # non-distributed training
            multiplier = 1 if num_gpu == 0 else num_gpu
            batch_size = dataset_opt['batch_size_per_gpu'] * multiplier
            num_workers = dataset_opt['num_worker_per_gpu'] * multiplier
        dataloader_args = dict(
            dataset=dataset,
            batch_size=batch_size,
            shuffle=False,
            num_workers=num_workers,
            sampler=sampler,
            drop_last=True)
        if sampler is None:
            dataloader_args['shuffle'] = True
        dataloader_args['worker_init_fn'] = partial(
            worker_init_fn, num_workers=num_workers, rank=rank, seed=seed) if seed is not None else None
    elif phase in ['val', 'test']:  # validation
        dataloader_args = dict(dataset=dataset, batch_size=1, shuffle=False, num_workers=0)
    else:
        raise ValueError(f"Wrong dataset phase: {phase}. Supported ones are 'train', 'val' and 'test'.")

    dataloader_args['pin_memory'] = dataset_opt.get('pin_memory', False)
    dataloader_args['persistent_workers'] = dataset_opt.get('persistent_workers', False)

    prefetch_mode = dataset_opt.get('prefetch_mode')
    if prefetch_mode == 'cpu':  # CPUPrefetcher
        num_prefetch_queue = dataset_opt.get('num_prefetch_queue', 1)
        logger = get_root_logger()
        logger.info(f'Use {prefetch_mode} prefetch dataloader: num_prefetch_queue = {num_prefetch_queue}')
        return PrefetchDataLoader(num_prefetch_queue=num_prefetch_queue, **dataloader_args)
    else:
        # prefetch_mode=None: Normal dataloader
        # prefetch_mode='cuda': dataloader for CUDAPrefetcher
        return torch.utils.data.DataLoader(**dataloader_args)
```

--------------------------------

### Get Initial Learning Rates

Source: https://basicsr.readthedocs.io/en/latest/_modules/basicsr/models/base_model

Retrieves the initial learning rates set for each parameter group across all optimizers. This is essential for implementing learning rate warm-up strategies.

```python
def _get_init_lr(self):
    """Get the initial lr, which is set by the scheduler.
    """
    init_lr_groups_l = []
    for optimizer in self.optimizers:
        init_lr_groups_l.append([v['initial_lr'] for v in optimizer.param_groups])
    return init_lr_groups_l
```

=== COMPLETE CONTENT === This response contains all available snippets from this library. No additional content exists. Do not make further requests.