TextBrewer

https://textbrewer.readthedocs.io/en/latest/
Admin
TextBrewer is a PyTorch-based model distillation toolkit for natural language processing, offering...
Tokens:42,069
Snippets:278
Trust Score:9.6
Update:4 months ago
Show doc for...
Tokens:
Raw
### Install TextBrewer from Github Source

Source: https://textbrewer.readthedocs.io/en/latest/Tutorial.html

Installs TextBrewer by cloning the GitHub repository and then installing locally. Requires Git.

```bash
git clone https://github.com/airaria/TextBrewer.git
pip install ./textbrewer

```

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

### SimpleAdaptor Function Example

Source: https://textbrewer.readthedocs.io/en/latest/Concepts.html

This example demonstrates how to implement an adaptor function to format model outputs for the TextBrewer distiller. It shows how to map model outputs like logits and hidden states to the expected dictionary keys.

```python
'''
Suppose the model outputs are: logits, sequence_output, total_loss
class MyModel():
  def forward(self, input_ids, attention_mask, labels, ...):
    ...
    return logits, sequence_output, total_loss

logits: Tensor of shape (batch_size, num_classes)
sequence_output: List of tensors of (batch_size, length, hidden_dim)
total_loss: scalar tensor

model inputs are:
input_ids      = batch[0] : input_ids (batch_size, length)
attention_mask = batch[1] : attention_mask (batch_size, length)
labels         = batch[2] : labels (batch_size, num_classes)
'''
def SimpleAdaptor(batch, model_outputs):
  return {'logits': (model_outputs[0],),
      'hidden': model.outputs[1],
      'inputs_mask': batch[1]}
```

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

### Install TextBrewer from PyPI

Source: https://textbrewer.readthedocs.io/en/latest/Tutorial.html

Installs the TextBrewer library using pip. Ensure Python and PyTorch are installed.

```bash
pip install textbrewer

```

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

### Quickstart: Distilling BERT with TextBrewer

Source: https://textbrewer.readthedocs.io/en/latest/Tutorial.html

Demonstrates distilling a BERT-base teacher model to a 3-layer BERT student model using TextBrewer's GeneralDistiller. Requires pre-trained teacher model, student model, dataloader, optimizer, and scheduler configurations.

```python
import textbrewer
from textbrewer import GeneralDistiller
from textbrewer import TrainingConfig, DistillationConfig

# Show the statistics of model parameters
print("\nteacher_model's parametrers:")
result, _ = textbrewer.utils.display_parameters(teacher_model,max_level=3)
print (result)

print("student_model's parametrers:")
result, _ = textbrewer.utils.display_parameters(student_model,max_level=3)
print (result)

# Define an adaptor for interpreting the model inputs and outputs
def simple_adaptor(batch, model_outputs):
    # The second and third elements of model outputs are the logits and hidden states
    return {'logits': model_outputs[1],
            'hidden': model_outputs[2]}

# Training configuration
train_config = TrainingConfig()
# Distillation configuration
# Matching different layers of the student and the teacher
# We match 0-0 and 8-2 here for demonstration
distill_config = DistillationConfig(
    intermediate_matches=[
    {'layer_T':0, 'layer_S':0, 'feature':'hidden', 'loss': 'hidden_mse','weight' : 1},
    {'layer_T':8, 'layer_S':2, 'feature':'hidden', 'loss': 'hidden_mse','weight' : 1}])

# Build distiller
distiller = GeneralDistiller(
    train_config=train_config, distill_config = distill_config,
    model_T = teacher_model, model_S = student_model,
    adaptor_T = simple_adaptor, adaptor_S = simple_adaptor)

# Start!
with distiller:
    distiller.train(optimizer, dataloader, num_epochs=1, scheduler_class=scheduler_class, scheduler_args=scheduler_args, callback=None)

```

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

### MultiTaskDistiller Training Loop Setup

Source: https://textbrewer.readthedocs.io/en/latest/_modules/textbrewer/distiller_multitask.html

Sets up the training loop for the MultiTaskDistiller, including optimizer and scheduler initialization, checkpointing strategy, and data iteration. It handles dataset size variations for weighted sampling.

```python
def train(self, optimizer, dataloaders, num_steps, scheduler_class=None, scheduler_args=None, scheduler=None, max_grad_norm = -1.0, tau=1, callback=None, batch_postprocessors=None, **args):
        """
        trains the student model.

        Args:
            optimizer: optimizer.
            dataloaders (dict): dict of dataset iterator. Keys are tasknames, values are corresponding dataloaders.
            num_steps (int): number of training steps.
            scheduler_class (class): the class of the scheduler to be constructed.
            scheduler_args (dict): arguments (excluding `optimizer`) passed to the `scheduler_class` to construct the scheduler object.
            scheduler (deprecated): used to adjust learning rate, optional, can be None, is deprecated in favor of `scheduler_class` and `scheduler_args`.
            max_grad_norm (float): Maximum norm for the gradients (-1 means no clipping). Default: -1.0
            tau (float): the probability of sampling an example from task `d` is proportional to ||d||^tau, where ||d|| is the size of `d`'s training set. If the size of any dataset is unknown, ignores tau and samples examples unifromly from each dataset.
            callback (Callable): function called after each epoch, can be None. It is called as ``callback(model=self.model_S, step = global_step)``. It can be used to do evaluation of the model at each checkpoint.
            batch_postprocessors (dict): a dict of batch_postprocessors. Keys are tasknames, values are corresponding batch_postprocessors. Each batch_postprocessor should take a batch and return a batch.
            **args: additional arguments fed to the model.
        """
        optimizer, scheduler, tqdm_disable = self.initialize_training(optimizer, scheduler_class, scheduler_args, scheduler)

        total_global_steps = num_steps
        ckpt_steps = int(self.t_config.ckpt_steps)
        num_steps = int(num_steps)
        print_every = ckpt_steps // self.print_freq
        if print_every == 0:
            print_every = ckpt_steps
        checkpoints = [ i * ckpt_steps for i in range(1,num_steps//ckpt_steps+1)] # + [total_global_steps]
        if checkpoints[-1] != total_global_steps:
            checkpoints.append(total_global_steps)
        logger.info(f"Total training steps: {total_global_steps}")
        logger.info(f"Checkpoints(step): {checkpoints}")

        dataiters = {k:cycle(v) for k,v in dataloaders.items()}
        if all(hasattr(v,'__len__') for v in dataloaders.values()):
            dataloader_sizes = {k:len(v) for k,v in dataloaders.items()}
            total_size = sum(v for k,v in dataloader_sizes.items())//self.t_config.gradient_accumulation_steps
            logger.info(f"Total size of all datasets (in number of batch_size):{total_size}")
            Z = sum(pow(v,tau) for v in dataloader_sizes.values())
            tasknames, sampling_weights = zip(*((k,pow(v,tau)/Z) for k,v in dataloader_sizes.items()))
        else:
            logger.info("The size of some datasets are unknown, so tau=1")
            tasknames = tuple(dataloaders.keys())
            sampling_weights = None


        global_step = 0
        writer_step = 0
        optimizer.zero_grad()
        while global_step < num_steps:
            global_step += 1
            for _ in range(self.t_config.gradient_accumulation_steps):

```

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

### Distillation Configuration Example

Source: https://textbrewer.readthedocs.io/en/latest/Experiments.html

Defines a distillation configuration with a specified temperature and intermediate layer matches. Other arguments default to their standard values.

```python
distill_config = DistillationConfig(temperature = 8, intermediate_matches = matches)
# Others arguments take the default values

```

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

### Training with Learning Rate Scheduler

Source: https://textbrewer.readthedocs.io/en/latest/Distillers.html

Trains the student model using a specified optimizer and dataloader. This example demonstrates how to configure a learning rate scheduler using `scheduler_class` and `scheduler_args`.

```python
from transformers import get_linear_schedule_with_warmup
distiller.train(optimizer, scheduler_class = get_linear_schedule_with_warmup, scheduler_args= {'num_warmup_steps': 100, 'num_training_steps': 1000})
```

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

### DataParallel Setup

Source: https://textbrewer.readthedocs.io/en/latest/_modules/textbrewer/distiller_basic.html

Configures models for data parallel training using PyTorch's DataParallel. This is used when `self.t_config.data_parallel` is true.

```python
self.model_S = torch.nn.DataParallel(self.model_S)
```

```python
self.model_T = [torch.nn.DataParallel(model_t) for model_t in self.model_T]
```

```python
self.model_T = {k:torch.nn.DataParallel(v) for k,v in self.model_T.items()}
```

```python
self.model_T = torch.nn.DataParallel(self.model_T)
```

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

### Initialize TrainingConfig with basic parameters

Source: https://textbrewer.readthedocs.io/en/latest/Configurations.html

Instantiate TrainingConfig by specifying the log directory and output directory. Other parameters will use their default values.

```python
train_config = TrainingConfig(log_dir=my_log_dir, output_dir=my_output_dir)
```

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

### DistributedDataParallel Setup

Source: https://textbrewer.readthedocs.io/en/latest/_modules/textbrewer/distiller_basic.html

Configures models for distributed training using PyTorch's DistributedDataParallel. This is used when `self.t_config.distributed` is true.

```python
self.model_T = {k:torch.nn.parallel.DistributedDataParallel(v, 
                        device_ids = [self.local_rank], output_device = self.local_rank,
                        find_unused_parameters = True) for k,v in self.model_T.items()}
```

```python
self.model_T = torch.nn.parallel.DistributedDataParallel(self.model_T,
                        device_ids = [self.local_rank], output_device = self.local_rank,
                        find_unused_parameters = True)
```

```python
self.projs[i] = torch.nn.parallel.DistributedDataParallel(proj,
                            device_ids = [self.local_rank], output_device = self.local_rank)
```

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

### GeneralDistiller Initialization

Source: https://textbrewer.readthedocs.io/en/latest/Distillers.html

Initializes the GeneralDistiller for single-teacher distillation with intermediate feature matching. Requires training and distillation configurations, teacher and student models, and their respective adaptors.

```python
GeneralDistiller(train_config, distill_config, model_T, model_S, adaptor_T, adaptor_S, custom_matches=None)
```

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

### BasicDistiller Initialization

Source: https://textbrewer.readthedocs.io/en/latest/Distillers.html

Initializes a BasicDistiller for single-teacher single-task distillation. Requires training and distillation configurations, teacher and student models, and their respective adaptors.

```python
distiller = BasicDistiller(train_config, distill_config, model_T, model_S, adaptor_T, adaptor_S)
```

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

### Adding a Custom L1 Loss Function

Source: https://textbrewer.readthedocs.io/en/latest/_modules/textbrewer/presets.html

Example of defining and registering a new Mean Absolute Error (MAE) loss function. This demonstrates how to create a custom loss and add it to the MATCH_LOSS_MAP.

```python
def my_L1_loss(feature_S, feature_T, mask=None):
    return (feature_S-feature_T).abs().mean()

MATCH_LOSS_MAP['my_L1_loss'] = my_L1_loss
```

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

### from_dict

Source: https://textbrewer.readthedocs.io/en/latest/Configurations.html

Construct configurations from a dictionary object.

```APIDOC
## from_dict

### Description
Construct configurations from a dictionary object.

### Class Method Signature
`from_dict(dict_object)`

### Parameters
- **dict_object** (dict) - A dictionary containing configuration parameters.
```

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

### GeneralDistiller Class Initialization

Source: https://textbrewer.readthedocs.io/en/latest/_modules/textbrewer/distiller_general.html

Initializes the GeneralDistiller with training and distillation configurations, teacher and student models, and adaptors. It sets up intermediate projection layers and custom match handlers if provided.

```python
from .distiller_utils import *
from .distiller_basic import BasicDistiller


class GeneralDistiller(BasicDistiller):
    """
    Supports intermediate features matching. **Recommended for single-teacher single-task distillation**.

    Args:
        train_config (:class:`TrainingConfig`): training configuration.
        distill_config (:class:`DistillationConfig`): distillation configuration.
        model_T (:class:`torch.nn.Module`): teacher model.
        model_S (:class:`torch.nn.Module`): student model.
        adaptor_T (Callable): teacher model's adaptor.
        adaptor_S (Callable): student model's adaptor.
        custom_matches (list): supports more flexible user-defined matches (testing).

    The roles of `adaptor_T` and `adaptor_S` are explained in :py:func:`adaptor`.

    """
    def __init__(self, train_config,
                 distill_config,
                 model_T,
                 model_S,
                 adaptor_T,
                 adaptor_S,
                 custom_matches: Optional[List[CustomMatch]] = None):
        # custom_matches=[{'module_T': module_T, 'module_S':module_S,
        #                 'loss': loss, 'weight': weight},...]
        super(GeneralDistiller, self).__init__(train_config, distill_config, model_T, model_S, adaptor_T, adaptor_S)

        self.projs = []
        self.projs_group = []
        for im in self.d_config.intermediate_matches:
            if im.proj is not None:
                projection = im.proj[0]
                dim_in = im.proj[1]
                dim_out = im.proj[2]
                self.projs_group.append(im.proj[3])
                self.projs.append(PROJ_MAP[projection](dim_in,dim_out))
                self.projs[-1].to(self.t_config.device)
            else:
                self.projs.append(None)
                self.projs_group.append(None)

        self.has_custom_matches = False
        if custom_matches:
            self.handles_T = []
            self.handles_S = []
            self.custom_matches_cache = {'hook_outputs_T': [], 'hook_outputs_S': [], 'match_proj_funcs': [],
                                         'match_weights':  [], 'match_losses':   [], 'match_proj_groups': []}
            for match in custom_matches:
                self.add_match(match)
            self.has_custom_matches = True
        
        self.d_config.is_caching_logits = False
```

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

### Long-Range Token Disorder Example

Source: https://textbrewer.readthedocs.io/en/latest/Utils.html

Performs long-range disordering by swapping halves of spans within a list. The `length` parameter controls the span size, accepting floats for relative length. Use for creating significant reordering in sequences.

```python
long_disorder([0,1,2,3,4,5,6,7,8,9,10], p=1, length=0.4)
```

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

### Initialize DistillationConfig with default values

Source: https://textbrewer.readthedocs.io/en/latest/_modules/textbrewer/configurations.html

Use this snippet to create a basic distillation configuration with default settings or to specify a custom temperature.

```python
from textbrewer import DistillationConfig

# simple configuration: use default values, or try different temperatures 
distill_config = DistillationConfig(temperature=8)
```

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

### Long-Range Disordering

Source: https://textbrewer.readthedocs.io/en/latest/_modules/textbrewer/data_utils.html

Performs long-range disordering by swapping halves of spans of a specified length. The length can be an absolute integer or a relative float. Example provided shows swapping halves of spans of length 4 (0.4 of 10 tokens).

```python
import numpy as np
import random


def long_disorder(tokens,p = 0.1, length=20):
    """
    Performs a long-range disordering. If ``length>1``, then swaps the two halves of each span of length `length` in `tokens`; 
    if ``length<=1``, treats `length` as the relative length. For example::
    
        >>>long_disorder([0,1,2,3,4,5,6,7,8,9,10], p=1, length=0.4)
        [2, 3, 0, 1, 6, 7, 4, 5, 8, 9]

    Args:
        tokens (list): list of tokens or token ids.
        p (list): probability to swaps the two halves of a spans at possible positions.
        length (int or float): length of the disordered span.
    Returns:
        a new disordered list
    """
    outputs = tokens[:]
    if int(length) <= 1:
        length = len(tokens)*length
    length = (int(length)+1) //2 * 2
    i = 0
    while i<=len(outputs)-length:
        if np.random.rand() < p:
            outputs[i:i+length//2], outputs[i+length//2:i+length] = outputs[i+length//2:i+length], outputs[i:i+length//2]
            i += length
        else:
            i += 1
    return outputs
```

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

### Initialize DistillationConfig

Source: https://textbrewer.readthedocs.io/en/latest/Configurations.html

Initializes the DistillationConfig with a specified temperature. This is a basic configuration for distillation.

```python
from textbrewer import DistillationConfig

# simple configuration: use default values, or try different temperatures
distill_config = DistillationConfig(temperature=8)
```

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

### Custom Temperature Scheduler Parameters

Source: https://textbrewer.readthedocs.io/en/latest/Presets.html

Demonstrates how to specify parameters for custom temperature schedulers like 'flsw' and 'cwsm' when creating a DistillationConfig.

```python
distill_config = DistillationConfig(
    temperature_scheduler = ['flsw', 1, 2]  # beta=1, gamma=2
)

distill_config = DistillationConfig(
    temperature_scheduler = ['cwsm', 1] # beta = 1
)
```

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

### Initialize DistillationConfig with multiple intermediate feature matches and projection

Source: https://textbrewer.readthedocs.io/en/latest/_modules/textbrewer/configurations.html

Set up distillation with multiple intermediate feature matches, including projection layers to handle differing feature dimensions between teacher and student models. This is useful when teacher and student models have different hidden dimensions.

```python
distill_config = DistillationConfig(
    temperature = 8, 
    intermediate_matches = [ \
    {'layer_T':0,  'layer_S':0, 'feature':'hidden','loss': 'hidden_mse', 'weight' : 1,'proj':['linear',384,768]},
    {'layer_T':4,  'layer_S':1, 'feature':'hidden','loss': 'hidden_mse', 'weight' : 1,'proj':['linear',384,768]},
    {'layer_T':8,  'layer_S':2, 'feature':'hidden','loss': 'hidden_mse', 'weight' : 1,'proj':['linear',384,768]},
    {'layer_T':12, 'layer_S':3, 'feature':'hidden','loss': 'hidden_mse', 'weight' : 1,'proj':['linear',384,768]}]
)
```

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

### MultiTeacherDistiller Initialization

Source: https://textbrewer.readthedocs.io/en/latest/Distillers.html

Initializes the MultiTeacherDistiller for distilling multiple teacher models into a single student model. Does not support intermediate feature matching. Requires training and distillation configurations, a list of teacher models, a student model, and their adaptors.

```python
MultiTeacherDistiller(train_config, distill_config, model_T, model_S, adaptor_T, adaptor_S)
```

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

### Initializing Training Components

Source: https://textbrewer.readthedocs.io/en/latest/_modules/textbrewer/distiller_basic.html

Initializes the optimizer and scheduler, potentially adding parameter groups for projection layers. It also handles FP16 initialization using Apex and sets up DistributedDataParallel for the student and teacher models if using distributed training.

```python
def initialize_training(self, optimizer, scheduler_class, scheduler_args, scheduler):
        # update optimizer for projection layer (used in GeneralDistiller)
        if hasattr(self,'projs'):
            for proj,proj_group in zip(self.projs, self.projs_group):
                if proj is not None:
                    assert isinstance(proj,nn.Module)
                    optimizer.add_param_group({**{'params':proj.parameters()},**proj_group})

        if hasattr(self,'has_custom_matches') and self.has_custom_matches:
            for proj_func,proj_group in zip(self.custom_matches_cache['match_proj_funcs'],
                                            self.custom_matches_cache['match_proj_groups']):
                if isinstance(proj_func,nn.Module):
                    optimizer.add_param_group({**{'params':proj_func.parameters()},**proj_group})

        logger.debug("Optimizer param group: ")
        logger.debug(f"{[[s.shape for s in g['params']] for g in optimizer.param_groups]}")

        # update scheduler
        if scheduler_class is not None:
            # overwrite scheduler
            scheduler = scheduler_class(**{'optimizer':optimizer},**scheduler_args)

        if self.t_config.fp16:
            if not has_apex:
                raise ImportError("Please install apex from https://www.github.com/nvidia/apex to use fp16 training.")
            if isinstance(self.model_T,(list,tuple)):
                models = [self.model_S] + list(self.model_T)
                models, optimizer = amp.initialize(models, optimizer, opt_level=self.t_config.fp16_opt_level)
                self.model_S = models[0]
                self.model_T =models[1:]
            elif isinstance(self.model_T,dict):
                tasknames, model_Ts = zip(*self.model_T.items())
                models = [self.model_S] + list(model_Ts)
                models, optimizer = amp.initialize(models, optimizer, opt_level=self.t_config.fp16_opt_level)
                self.model_S = models[0]
                self.model_T = dict(zip(tasknames,models[1:]))
            else:
                (self.model_S, self.model_T), optimizer = amp.initialize([self.model_S, self.model_T], optimizer, opt_level=self.t_config.fp16_opt_level)
        if self.local_rank != -1:
            self.model_S = torch.nn.parallel.DistributedDataParallel(self.model_S,
                        device_ids = [self.local_rank], output_device = self.local_rank,
                        find_unused_parameters = True)
            if isinstance(self.model_T,(list,tuple)):
                self.model_T = [torch.nn.parallel.DistributedDataParallel(model_t,

```

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

### BasicTrainer Initialization

Source: https://textbrewer.readthedocs.io/en/latest/_modules/textbrewer/distiller_train.html

Initializes the BasicTrainer class with training configuration, model, and an adaptor. Sets up logging and distributed training if applicable.

```python
def __init__(self, train_config: TrainingConfig,
                 model: torch.nn.Module, adaptor):
        super(BasicTrainer, self).__init__()
        self.t_config = train_config
        self.model = model
        self.adaptor = adaptor
        self.local_rank = self.t_config.local_rank
        self.rank = 0
        if self.local_rank != -1:
            self.rank = torch.distributed.get_rank()
        if self.t_config.log_dir is not None and self.rank == 0:
            self.tb_writer = SummaryWriter(log_dir = self.t_config.log_dir)
        else:
            self.tb_writer = no_op
        self.print_freq = 20
```

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

### BasicDistiller

Source: https://textbrewer.readthedocs.io/en/latest/Distillers.html

Initializes a BasicDistiller object for single-teacher single-task distillation.

```APIDOC
## BasicDistiller

### Description
Performs **single-teacher single-task** distillation, provides basic distillation strategies.

### Parameters
* **train_config** (`TrainingConfig`) – training configuration.
* **distill_config** (`DistillationConfig`) – distillation configuration.
* **model_T** (`torch.nn.Module`) – teacher model.
* **model_S** (`torch.nn.Module`) – student model.
* **adaptor_T** (_Callable_) – teacher model’s adaptor.
* **adaptor_S** (_Callable_) – student model’s adaptor.

The roles of adaptor_T and adaptor_S are explained in `adaptor()`.
```

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

### from_json_file

Source: https://textbrewer.readthedocs.io/en/latest/Configurations.html

Construct configurations from a JSON file.

```APIDOC
## from_json_file

### Description
Construct configurations from a JSON file.

### Class Method Signature
`from_json_file(json_filename)`

### Parameters
- **json_filename** (str) - The path to the JSON configuration file.
```

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

### Using Custom Modules in DistillationConfig

Source: https://textbrewer.readthedocs.io/en/latest/Presets.html

Illustrates how to reference custom-defined loss functions and weight schedulers within the DistillationConfig.

```python
distill_config = DistillationConfig(
    kd_loss_weight_scheduler = 'my_weight_scheduler'
    intermediate_matches = [{'layer_T':0, 'layer_S':0, 'feature':'hidden','loss': 'my_L1_loss', 'weight' : 1}]
...)
```

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

### Configurations

Source: https://textbrewer.readthedocs.io/en/latest/index.html

Configuration classes for setting up training and distillation processes.

```APIDOC
## Configurations

### TrainingConfig

**Description**: Configuration for training parameters.

### DistillationConfig

**Description**: Configuration for distillation specific parameters.
```

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

### MultiTaskDistiller

Source: https://textbrewer.readthedocs.io/en/latest/_modules/textbrewer/distiller_multitask.html

Initializes the MultiTaskDistiller. It takes training and distillation configurations, teacher and student models, and their corresponding adaptors.

```APIDOC
## MultiTaskDistiller

### Description
Initializes the MultiTaskDistiller. It takes training and distillation configurations, teacher and student models, and their corresponding adaptors.

### Parameters
#### `train_config` (:class:`TrainingConfig`)
- Training configuration object.
#### `distill_config` (:class:`DistillationConfig`)
- Distillation configuration object.
#### `model_T` (dict)
- Dictionary of teacher models, where keys are task names and values are the corresponding models.
#### `model_S` (torch.nn.Module)
- The student model.
#### `adaptor_T` (dict)
- Dictionary of teacher adaptors, where keys are task names and values are the corresponding adaptors.
#### `adaptor_S` (dict)
- Dictionary of student adaptors, where keys are task names and values are the corresponding adaptors.

```

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

### Adding Custom Loss and Scheduler to Presets

Source: https://textbrewer.readthedocs.io/en/latest/Presets.html

Shows how to extend the MATCH_LOSS_MAP and WEIGHT_SCHEDULER dictionaries with custom functions.

```python
MATCH_LOSS_MAP['my_L1_loss'] = my_L1_loss
WEIGHT_SCHEDULER['my_weight_scheduler'] = my_weight_scheduler
```

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

### MultiTaskDistiller Initialization

Source: https://textbrewer.readthedocs.io/en/latest/_modules/textbrewer/distiller_multitask.html

Initializes the MultiTaskDistiller with training and distillation configurations, teacher and student models, and their respective adaptors. It enforces consistency in the number of adaptors and models for each task and disables caching of logits.

```python
from .distiller_utils import *
from .distiller_general import GeneralDistiller


class MultiTaskDistiller(GeneralDistiller):
    """
    distills multiple teacher models (of different tasks) into a single student. **It supports intermediate feature matching since 0.2.1**.

    Args:
        train_config (:class:`TrainingConfig`): training configuration.
        distill_config (:class:`DistillationConfig`): distillation configuration.
        model_T (dict): dict of teacher models: {task1:model1, task2:model2, .... }. Keys are tasknames.
        model_S (torch.nn.Module): student model.
        adaptor_T (dict): dict of teacher adaptors: {task1:adpt1, task2:adpt2, .... }. Keys are tasknames.
        adaptor_S (dict): dict of student adaptors: {task1:adpt1, task2:adpt2, .... }. Keys are tasknames.

    """

    def __init__(self, train_config,
                 distill_config,
                 model_T,
                 model_S,
                 adaptor_T,
                 adaptor_S):

        super(MultiTaskDistiller, self).__init__(
            train_config, distill_config,
            model_T, model_S,
            adaptor_T, adaptor_S)
        if hasattr(self.adaptor_T,'__iter__'):
            assert len(self.adaptor_T)==len(self.model_T)==len(self.adaptor_S)
        #assert (self.d_config.kd_loss_weight_scheduler is None) and (self.d_config.hard_label_weight_scheduler is None),
        #        "MultiTaskDistiller does not support WEIGHT_SCHEDULER in the current version."

        self.d_config.is_caching_logits = False
```

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

### BasicTrainer Initialization

Source: https://textbrewer.readthedocs.io/en/latest/Distillers.html

Initializes the BasicTrainer for supervised training. Used for training teacher models.

```python
trainer = BasicTrainer(train_config, model, adaptor)
```

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

### GeneralDistiller Train Method

Source: https://textbrewer.readthedocs.io/en/latest/Distillers.html

Trains the student model using the GeneralDistiller. Supports specifying number of epochs or steps, learning rate schedulers via class and arguments, gradient clipping, and callbacks.

```python
distiller.train(optimizer, dataloader, num_epochs=None, scheduler_class=None, scheduler_args=None, scheduler=None, max_grad_norm=-1.0, num_steps=None, callback=None, batch_postprocessor=None, **args)
```

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

### TEMPERATURE_SCHEDULER Configuration

Source: https://textbrewer.readthedocs.io/en/latest/_modules/textbrewer/presets.html

A DynamicKeyDict for selecting temperature schedulers. Supports 'constant', 'flsw', and 'cwsm' with specific parameter requirements.

```python
TEMPERATURE_SCHEDULER=DynamicKeyDict(
    {'constant': constant_temperature_scheduler,
     'flsw': flsw_temperature_scheduler_builder,
     'cwsm':cwsm_temperature_scheduler_builder})

```

```python
#flsw
distill_config = DistillationConfig(
    temperature_scheduler = ['flsw', 1， 2]  # beta=1, gamma=2
)

#cwsm
distill_config = DistillationConfig(
    temperature_scheduler = ['cwsm', 1] # beta = 1
)
```

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

### GeneralDistiller

Source: https://textbrewer.readthedocs.io/en/latest/Distillers.html

Supports intermediate features matching. Recommended for single-teacher single-task distillation. Initializes the distiller with training and distillation configurations, teacher and student models, and their respective adaptors.

```APIDOC
## GeneralDistiller

### Description
Initializes the GeneralDistiller with training and distillation configurations, teacher and student models, and their respective adaptors. Supports intermediate features matching and is recommended for single-teacher single-task distillation.

### Parameters
* **train_config** (TrainingConfig) – training configuration.
* **distill_config** (DistillationConfig) – distillation configuration.
* **model_T** (torch.nn.Module) – teacher model.
* **model_S** (torch.nn.Module) – student model.
* **adaptor_T** (Callable) – teacher model’s adaptor.
* **adaptor_S** (Callable) – student model’s adaptor.
* **custom_matches** (Optional[List[textbrewer.distiller_utils.CustomMatch]]) – supports more flexible user-defined matches (testing).
```

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

### MultiTaskDistiller Initialization

Source: https://textbrewer.readthedocs.io/en/latest/Distillers.html

Initializes the MultiTaskDistiller for multi-task distillation. Supports intermediate feature matching.

```python
distiller = MultiTaskDistiller(train_config, distill_config, model_T, model_S, adaptor_T, adaptor_S)
```

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

### Basic Training Loop with Gradient Accumulation

Source: https://textbrewer.readthedocs.io/en/latest/_modules/textbrewer/distiller_basic.html

This snippet illustrates the main training loop of the basic distiller. It handles epoch iteration, batch processing, loss calculation, gradient accumulation, and checkpointing. Use this for standard training scenarios.

```python
        if print_every == 0:
            print_every = train_steps_per_epoch
        checkpoints = [int(train_steps_per_epoch*ci/self.t_config.ckpt_frequency) for ci in range(self.t_config.ckpt_frequency)]
        logger.info(f"Training steps per epoch: {train_steps_per_epoch}")
        logger.info(f"Checkpoints(step): {checkpoints}")

        global_step = 0
        writer_step = 0

        if self.d_config.is_caching_logits is True:
            logger.info(f"Caching batches and teacher's logits...")
            for step, batch in tqdm(enumerate(dataloader),disable=tqdm_disable):
                self.cache_logits(batch, args, batch_postprocessor)

        for current_epoch in tqdm(range(int(num_epochs)),disable=tqdm_disable):
            if self.local_rank != -1 and hasattr(dataloader,'sampler'):
                dataloader.sampler.set_epoch(current_epoch)  #In distributed mode, calling the set_epoch method is needed to make shuffling work;
            logger.info(f"Epoch {current_epoch+1}")
            optimizer.zero_grad()
            if self.d_config.is_caching_logits:
                random.shuffle(self.logits_cache)
                dataloader = self.logits_cache
            logger.info(f"Length of current epoch in forward batch: {len(dataloader)}")
            for step, batch in tqdm(enumerate(dataloader),disable=tqdm_disable):
                if self.d_config.is_caching_logits is False and batch_postprocessor is not None:
                        batch = batch_postprocessor(batch)
                total_loss, losses_dict = self.train_on_batch(batch,args)

                self.write_loss(total_loss, writer_step, losses_dict)
                writer_step += 1

                total_loss /= self.t_config.gradient_accumulation_steps
                if self.t_config.fp16:
                    with amp.scale_loss(total_loss,optimizer) as scaled_loss:
                        scaled_loss.backward()
                else:
                    total_loss.backward()

                if (step+1)%self.t_config.gradient_accumulation_steps == 0:
                    if max_grad_norm > 0:
                        if self.t_config.fp16:
                            torch.nn.utils.clip_grad_norm_(amp.master_params(optimizer), max_grad_norm)
                        else:
                            torch.nn.utils.clip_grad_norm_(self.model_S.parameters(), max_grad_norm) 
                    optimizer.step()
                    if scheduler is not None:
                        scheduler.step()
                    optimizer.zero_grad()
                    global_step += 1
                    if self.d_config.kd_loss_weight_scheduler is not None:
                        self.d_config.kd_loss_weight = \
                            self.d_config.kd_loss_weight_scheduler(global_step/total_global_steps)
                    if self.d_config.hard_label_weight_scheduler is not None:
                        self.d_config.hard_label_weight = \
                            self.d_config.hard_label_weight_scheduler(global_step/total_global_steps)

                    if (global_step) % print_every == 0:
                        logger.info(f"Global step: {global_step}, epoch step:{step+1}")
                    if (global_step%train_steps_per_epoch in checkpoints) \
                            and ((current_epoch+1)%self.t_config.ckpt_epoch_frequency==0 or current_epoch+1==num_epochs):
                        self.save_and_callback(global_step, step, current_epoch, callback)

            logger.info(f"Epoch {current_epoch+1} finished")
```

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

### BasicTrainer

Source: https://textbrewer.readthedocs.io/en/latest/Distillers.html

Initializes a BasicTrainer for performing supervised training, not distillation. Can be used for training the teacher model.

```APIDOC
## class textbrewer.BasicTrainer

### Description
Performs supervised training, not distillation. It can be used for training the teacher model.

### Parameters
* **train_config** (TrainingConfig) - training configuration.
* **model** (torch.nn.Module) - model to be trained.
* **adaptor** (Callable) - The role of adaptor is explained in `adaptor()`.
```

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

### MultiTaskDistiller Training Method

Source: https://textbrewer.readthedocs.io/en/latest/Distillers.html

Trains the student model using the MultiTaskDistiller. Allows for custom optimizers, dataloaders, schedulers, and callbacks.

```python
distiller.train(optimizer, dataloaders, num_steps)
```

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

### MultiTeacherDistiller.train

Source: https://textbrewer.readthedocs.io/en/latest/Distillers.html

Trains the student model using the provided optimizer, scheduler, dataloader, and training configurations. Allows for custom callbacks and batch post-processing. See BasicDistiller.train().

```APIDOC
## MultiTeacherDistiller.train

### Description
Trains the student model. This method is similar to `GeneralDistiller.train()` and is described in `BasicDistiller.train()`.

### Method
`train`

### Parameters
* **optimizer** – optimizer.
* **scheduler** – scheduler.
* **dataloader** – dataset iterator.
* **num_epochs** – number of training epochs.
* **num_steps** (int) – number of training steps. If it is not None, distiller will ignore num_epochs and trains for num_steps, and dataloader can have an unkonwn size, i.e., has no __len__ attribute. Dataloader will be cycled automatically after iterating over the whole dataset.
* **callback** (Callable) – function called after each epoch, can be None. It is called as `callback(model=self.model_S, step = global_step)`. It can be used to evaluate the model at each checkpoint.
* **batch_postprocessor** (Callable) – a function for post-processing batches. It should take a batch and return a batch. Its output is fed to the models and adaptors.
* **args** – additional arguments fed to the model.
```

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

### GeneralDistiller.train

Source: https://textbrewer.readthedocs.io/en/latest/Distillers.html

Trains the student model using the provided optimizer, dataloader, and training configurations. Allows for custom callbacks and batch post-processing.

```APIDOC
## GeneralDistiller.train

### Description
Trains the student model. This method takes an optimizer, dataloader, and optional training parameters like number of epochs or steps, learning rate scheduler configuration, gradient clipping, and callbacks.

### Method
`train`

### Parameters
* **optimizer** – optimizer.
* **dataloader** – dataset iterator.
* **num_epochs** (int) – number of training epochs.
* **num_steps** (int) – number of training steps. If it is not None, distiller will ignore num_epochs and trains for num_steps, and dataloader can have an unkonwn size, i.e., has no __len__ attribute. Dataloader will be cycled automatically after iterating over the whole dataset.
* **callback** (Callable) – function called after each epoch, can be None. It is called as `callback(model=self.model_S, step = global_step)`. It can be used to evaluate the model at each checkpoint.
* **batch_postprocessor** (Callable) – a function for post-processing batches. It should take a batch and return a batch. Its output is fed to the models and adaptors.
* **scheduler_class** (class) – the class of the scheduler to be constructed.
* **scheduler_args** (dict) – arguments (excluding optimizer) passed to the scheduler_class to construct the scheduler object.
* **scheduler** (deprecated) – used to adjust learning rate, optional, can be None, is deprecated in favor of scheduler_class and scheduler_args.
* **max_grad_norm** (float) – Maximum norm for the gradients (-1 means no clipping). Default: -1.0
* **args** – additional arguments fed to the model.

### Notes
* If the batch is a list or tuple, model is called as: `model(*batch, **args)`. Make sure the order of elements in the batch matches their order in `model.forward`.
* If the batch is a dict, model is called as: `model(**batch,**args)`. Make sure the keys of the batch match the arguments of the `model.forward`.
* If you want to provide a lr scheduler, DON’T USE scheduler , use scheduler_class and scheduler_args instead. Example:
```python
from transformers import get_linear_schedule_with_warmup
distiller.train(optimizer, scheduler_class = get_linear_schedule_with_warmup, scheduler_args= {'num_warmup_steps': 100, 'num_training_steps': 1000})
```
```

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

### BasicDistiller.train

Source: https://textbrewer.readthedocs.io/en/latest/Distillers.html

Trains the student model using the provided optimizer and dataloader.

```APIDOC
## train

### Description
Trains the student model.

### Parameters
* **optimizer** – optimizer.
* **dataloader** – dataset iterator.
* **num_epochs** (_int_) – number of training epochs.
* **num_steps** (_int_) – number of training steps. If it is not None, distiller will ignore num_epochs and trains for num_steps, and dataloader can have an unkonwn size, i.e., has no __len__ attribute. Dataloader will be cycled automatically after iterating over the whole dataset.
* **callback** (_Callable_) – function called after each epoch, can be None. It is called as `callback(model=self.model_S, step = global_step)`. It can be used to evaluate the model at each checkpoint.
* **batch_postprocessor** (_Callable_) – a function for post-processing batches. It should take a batch and return a batch. Its output is fed to the models and adaptors.
* **scheduler_class** (_class_) – the class of the scheduler to be constructed.
* **scheduler_args** (_dict_) – arguments (excluding optimizer) passed to the scheduler_class to construct the scheduler object. See the example below.
* **scheduler** (_deprecated_) – used to adjust learning rate, optional, can be None, is deprecated in favor of scheduler_class and scheduler_args.
* **max_grad_norm** (_float_) – Maximum norm for the gradients (-1 means no clipping). Default: -1.0
* ****args** – additional arguments fed to the model.

### Notes
* If the batch is a list or tuple, model is called as: `model(*batch, **args)`. Make sure the order of elements in the batch matches their order in `model.forward`.
* If the batch is a dict, model is called as: `model(**batch,**args)`. Make sure the keys of the batch match the arguments of the `model.forward`.

### Example Usage for LR Scheduler
If you want to provide a lr scheduler, DON’T USE scheduler , use scheduler_class and scheduler_args instead. Example:
```python
from transformers import get_linear_schedule_with_warmup
distiller.train(optimizer, scheduler_class = get_linear_schedule_with_warmup, scheduler_args= {'num_warmup_steps': 100, 'num_training_steps': 1000})
```
```

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

### Build a weight initializer function

Source: https://textbrewer.readthedocs.io/en/latest/_modules/textbrewer/utils.html

Creates a function that initializes weights for Linear and Embedding layers, and biases for Linear layers. Use the returned function to apply these initializations to a PyTorch module.

```python
def initializer_builder(std):
    _std = std
    def init_weights(module):
        if isinstance(module, (torch.nn.Linear, torch.nn.Embedding)):
            module.weight.data.normal_(mean=0.0, std=_std)
        if isinstance(module, torch.nn.Linear) and module.bias is not None:
            module.bias.data.zero_()
    return init_weights
```

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

### Add Match by Module

Source: https://textbrewer.readthedocs.io/en/latest/_modules/textbrewer/distiller_general.html

Registers forward hooks for student and teacher modules, appends match details to cache, and moves projection functions to the appropriate device.

```python
def add_match_by_module(self,module_T : torch.nn.Module,
                                 module_S : torch.nn.Module,
                                 proj_func, proj_group,
                                 match_weight, match_loss):

        self.handles_T = module_T.register_forward_hook(self._hook_T)
        self.handles_S = module_S.register_forward_hook(self._hook_S)
        self.custom_matches_cache['match_weights'].append(match_weight)
        self.custom_matches_cache['match_losses'].append(match_loss)
        self.custom_matches_cache['match_proj_funcs'].append(proj_func)
        if isinstance(proj_func,nn.Module):
            self.custom_matches_cache['match_proj_funcs'][-1].to(self.t_config.device)
        self.custom_matches_cache['match_proj_groups'].append(proj_group)
```

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

### TrainingConfig

Source: https://textbrewer.readthedocs.io/en/latest/Configurations.html

Configurations related to general model training. This class allows users to specify various parameters to control the training process, such as checkpointing frequency, logging and output directories, device usage, and mixed precision settings.

```APIDOC
## TrainingConfig

### Description
Configurations related to general model training. This class allows users to specify various parameters to control the training process, such as checkpointing frequency, logging and output directories, device usage, and mixed precision settings.

### Class Signature
`textbrewer.TrainingConfig(gradient_accumulation_steps=1, ckpt_frequency=1, ckpt_epoch_frequency=1, ckpt_steps=None, log_dir=None, output_dir='./saved_models', device='cuda', fp16=False, fp16_opt_level='O1', data_parallel=False, local_rank=-1)`

### Parameters
#### Constructor Parameters
- **gradient_accumulation_steps** (int) - accumulates gradients before optimization to reduce GPU memory usage. It calls `optimizer.step()` every gradient_accumulation_steps backward steps.
- **ckpt_frequency** (int) - stores model weights ckpt_frequency times every epoch.
- **ckpt_epoch_frequency** (int) - stores model weights every ckpt_epoch_frequency epochs.
- **ckpt_steps** (int) - if `num_steps` is passed to `distiller.train()`, saves the model every **ckpt_steps**, meanwhile ignore ckpt_frequency and ckpt_epoch_frequency.
- **log_dir** (str) - directory to save the tensorboard log file. Set it to `None` to disable tensorboard.
- **output_dir** (str) - directory to save model weights.
- **device** (str or torch.device) - training on CPU or GPU.
- **fp16** (bool) - if `True`, enables mixed precision training using Apex.
- **fp16_opt_level** (str) - Pure or mixed precision optimization level. Accepted values are “O0”, “O1”, “O2”, and “O3”. See Apex documentation for details.
- **data_parallel** (bool) - If `True`, wraps the models with `torch.nn.DataParallel`.
- **local_rank** (int) - the local rank of the current processes. A non-negative value means that we are in the distributed training mode with `DistributedDataParallel`.

### Notes
- To perform data parallel (DP) training, you could either wrap the models with `torch.nn.DataParallel` outside TextBrewer by yourself, or leave the work for TextBrewer by setting **data_parallel** to `True`.
- To enable both data parallel training and mixed precision training, you should set **data_parallel** to `True`, and DO NOT wrap the models by yourself.
- In some experiments, we have observed slowing down in the speed with `torch.nn.DataParallel`.
- To perform distributed data parallel (DDP) training, you should call `torch.distributed.init_process_group` before initializing a TrainingConfig; and pass the **raw** (unwrapped) model when initializing the distiller.
- DP and DDP are mutual exclusive.

### Example
```python
# Usually just need to set log_dir and output_dir and leave others default
train_config = TrainingConfig(log_dir=my_log_dir, output_dir=my_output_dir)

# Stores the model at the end of each epoch
train_config = TrainingConfig(ckpt_frequency=1, ckpt_epoch_frequency=1)
# Stores the model twice (at the middle and at the end) in each epoch
train_config = TrainingConfig(ckpt_frequency=2, ckpt_epoch_frequency=1)
# Stores the model once every two epochs
train_config = TrainingConfig(ckpt_frequency=1, ckpt_epoch_frequency=2)
```
```

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

### MultiTeacherDistiller Train Method

Source: https://textbrewer.readthedocs.io/en/latest/Distillers.html

Trains the student model using the MultiTeacherDistiller. This method is similar to BasicDistiller.train and accepts optimizer, scheduler, dataloader, number of epochs, and optional arguments.

```python
distiller.train(self, optimizer, scheduler, dataloader, num_epochs, num_steps=None, callback=None, batch_postprocessor=None, **args)
```

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

### MultiTaskDistiller

Source: https://textbrewer.readthedocs.io/en/latest/Distillers.html

Initializes a MultiTaskDistiller for distilling multiple teacher models into a single student model. Supports intermediate feature matching.

```APIDOC
## class textbrewer.MultiTaskDistiller

### Description
Initializes a MultiTaskDistiller for distilling multiple teacher models (of different tasks) into a single student. It supports intermediate feature matching.

### Parameters
* **train_config** (TrainingConfig) - training configuration.
* **distill_config** (DistillationConfig) - distillation configuration.
* **model_T** (dict) - dict of teacher models: {task1:model1, task2:model2, …. }. Keys are tasknames.
* **model_S** (torch.nn.Module) - student model.
* **adaptor_T** (dict) - dict of teacher adaptors: {task1:adpt1, task2:adpt2, …. }. Keys are tasknames.
* **adaptor_S** (dict) - dict of student adaptors: {task1:adpt1, task2:adpt2, …. }. Keys are tasknames.
```
Tokens:
Raw