Repository: microsoft/nni
File: examples/tutorials/pruning_bert_glue.py (Python)

Code before:
    }

    return train_dataloader, validation_dataloaders


train_dataloader, validation_dataloaders = prepare_dataloaders()

# %%
# Training function & evaluation function.

import functools
import time

import torch.nn.functional as F
from datasets import load_metric
from transformers.modeling_outputs import SequenceClassifierOutput

from nni.common.version import torch_version_is_2


if torch_version_is_2():
    from torch.optim.lr_scheduler import LRScheduler
    SCHEDULER = LRScheduler
else:
    from torch.optim.lr_scheduler import _LRScheduler
    SCHEDULER = _LRScheduler


def training(model: torch.nn.Module,
             optimizer: torch.optim.Optimizer,
             criterion: Callable[[torch.Tensor, torch.Tensor], torch.Tensor],
             lr_scheduler: SCHEDULER = None, # type: ignore
             max_steps: int = None,
             max_epochs: int = None,
             train_dataloader: DataLoader = None,
             distillation: bool = False,
             teacher_model: torch.nn.Module = None,
             distil_func: Callable = None,
             log_path: str = Path(log_dir) / 'training.log',
             save_best_model: bool = False,
             save_path: str = None,
             evaluation_func: Callable = None,
             eval_per_steps: int = 1000,
             device=None):

    assert train_dataloader is not None

    model.train()
    if teacher_model is not None:
        teacher_model.eval()
    current_step = 0

Review comment (suggestion):
what about put this `if ... else ...` in some where and we just `from xxx import SCHEDULER` here, then we don't need to copy paste it everywhere.

Code after:
    }

    return train_dataloader, validation_dataloaders


train_dataloader, validation_dataloaders = prepare_dataloaders()

# %%
# Training function & evaluation function.

import functools
import time

import torch.nn.functional as F
from datasets import load_metric
from transformers.modeling_outputs import SequenceClassifierOutput

from nni.common.types import SCHEDULER


def training(model: torch.nn.Module,
             optimizer: torch.optim.Optimizer,
             criterion: Callable[[torch.Tensor, torch.Tensor], torch.Tensor],
             lr_scheduler: SCHEDULER = None,
             max_steps: int = None,
             max_epochs: int = None,
             train_dataloader: DataLoader = None,
             distillation: bool = False,
             teacher_model: torch.nn.Module = None,
             distil_func: Callable = None,
             log_path: str = Path(log_dir) / 'training.log',
             save_best_model: bool = False,
             save_path: str = None,
             evaluation_func: Callable = None,
             eval_per_steps: int = 1000,
             device=None):

    assert train_dataloader is not None

    model.train()
    if teacher_model is not None:
        teacher_model.eval()
    current_step = 0
    best_result = 0

    total_epochs = max_steps // len(train_dataloader) + 1 if max_steps else max_epochs if max_epochs else 3
    total_steps = max_steps if max_steps else total_epochs * len(train_dataloader)

    print(f'Training {total_epochs} epochs, {total_steps} steps...')

    for current_epoch in range(total_epochs):