AutoAndroidController/py/venv/Lib/site-packages/modelscope/trainers/audio/kws_farfield_trainer.py

import datetime
import glob
import math
import os
import pickle
from typing import Callable, Dict, Optional

import numpy as np
import torch
from torch import nn as nn
from torch import optim as optim

from modelscope.metainfo import Trainers
from modelscope.models import Model, TorchModel
from modelscope.msdatasets.dataset_cls.custom_datasets.audio import (
    KWSDataLoader, KWSDataset)
from modelscope.trainers.base import BaseTrainer
from modelscope.trainers.builder import TRAINERS
from modelscope.utils.audio.audio_utils import update_conf
from modelscope.utils.constant import DEFAULT_MODEL_REVISION, ModelFile
from modelscope.utils.data_utils import to_device
from modelscope.utils.device import create_device
from modelscope.utils.logger import get_logger
from modelscope.utils.torch_utils import (get_dist_info, get_local_rank,
                                          init_dist, is_master)

logger = get_logger()

BASETRAIN_CONF_EASY = 'basetrain_easy'
BASETRAIN_CONF_NORMAL = 'basetrain_normal'
BASETRAIN_CONF_HARD = 'basetrain_hard'
FINETUNE_CONF_EASY = 'finetune_easy'
FINETUNE_CONF_NORMAL = 'finetune_normal'
FINETUNE_CONF_HARD = 'finetune_hard'
CKPT_PREFIX = 'checkpoint'

EASY_RATIO = 0.1
NORMAL_RATIO = 0.6
HARD_RATIO = 0.3
BASETRAIN_RATIO = 0.5


@TRAINERS.register_module(module_name=Trainers.speech_dfsmn_kws_char_farfield)
class KWSFarfieldTrainer(BaseTrainer):
    DEFAULT_WORK_DIR = './work_dir'
    conf_keys = (BASETRAIN_CONF_EASY, FINETUNE_CONF_EASY,
                 BASETRAIN_CONF_NORMAL, FINETUNE_CONF_NORMAL,
                 BASETRAIN_CONF_HARD, FINETUNE_CONF_HARD)

    def __init__(self,
                 model: str,
                 work_dir: str,
                 cfg_file: Optional[str] = None,
                 arg_parse_fn: Optional[Callable] = None,
                 model_revision: Optional[str] = DEFAULT_MODEL_REVISION,
                 custom_conf: Optional[dict] = None,
                 **kwargs):

        if isinstance(model, str):
            self.model_dir = self.get_or_download_model_dir(
                model, model_revision)
            if cfg_file is None:
                cfg_file = os.path.join(self.model_dir,
                                        ModelFile.CONFIGURATION)
        else:
            assert cfg_file is not None, 'Config file should not be None if model is not from pretrained!'
            self.model_dir = os.path.dirname(cfg_file)

        super().__init__(cfg_file, arg_parse_fn)

        # the number of model output dimension
        # should update config outside the trainer, if user need more wake word
        num_syn = kwargs.get('num_syn', None)
        if num_syn:
            self.cfg.model.num_syn = num_syn
        self._num_classes = self.cfg.model.num_syn
        self.model = self.build_model()
        self.work_dir = work_dir

        if kwargs.get('launcher', None) is not None:
            init_dist(kwargs['launcher'])

        _, world_size = get_dist_info()
        self._dist = world_size > 1

        device_name = kwargs.get('device', 'gpu')
        if self._dist:
            local_rank = get_local_rank()
            device_name = f'cuda:{local_rank}'

        self.device = create_device(device_name)
        # model placement
        if self.device.type == 'cuda':
            self.model.to(self.device)

        if 'max_epochs' not in kwargs:
            assert hasattr(
                self.cfg.train, 'max_epochs'
            ), 'max_epochs is missing from the configuration file'
            self._max_epochs = self.cfg.train.max_epochs
        else:
            self._max_epochs = kwargs['max_epochs']
        self._train_iters = kwargs.get('train_iters_per_epoch', None)
        self._val_iters = kwargs.get('val_iters_per_epoch', None)
        if self._train_iters is None:
            self._train_iters = self.cfg.train.train_iters_per_epoch
        if self._val_iters is None:
            self._val_iters = self.cfg.evaluation.val_iters_per_epoch
        dataloader_config = self.cfg.train.dataloader
        self._threads = kwargs.get('workers', None)
        if self._threads is None:
            self._threads = dataloader_config.workers_per_gpu
        self._single_rate = BASETRAIN_RATIO
        if 'single_rate' in kwargs:
            self._single_rate = kwargs['single_rate']
        self._batch_size = dataloader_config.batch_size_per_gpu
        next_epoch = kwargs.get('next_epoch', 1)
        self._current_epoch = next_epoch - 1
        if 'model_bin' in kwargs:
            model_bin_file = os.path.join(self.model_dir, kwargs['model_bin'])
            self.model = torch.load(model_bin_file, weights_only=True)
        elif self._current_epoch > 0:
            # load checkpoint
            ckpt_file_pattern = os.path.join(
                self.work_dir, f'{CKPT_PREFIX}_{self._current_epoch:04d}*.pth')
            ckpt_files = glob.glob(ckpt_file_pattern)
            if len(ckpt_files) == 1:
                logger.info('Loading model from checkpoint: %s', ckpt_files[0])
                self.model = torch.load(ckpt_files[0], weights_only=True)
            elif len(ckpt_files) == 0:
                raise FileNotFoundError(
                    f'Failed to load checkpoint file like '
                    f'{ckpt_file_pattern}. File not found!')
            else:
                raise AssertionError(f'Expecting one but multiple checkpoint'
                                     f' files are found: {ckpt_files}')

        # build corresponding optimizer and loss function
        lr = self.cfg.train.optimizer.lr
        self.optimizer = optim.Adam(self.model.parameters(), lr)
        self.loss_fn = nn.CrossEntropyLoss()
        self.data_val = None
        self.json_log_path = os.path.join(self.work_dir,
                                          '{}.log.json'.format(self.timestamp))
        self.conf_files = []
        for conf_key in self.conf_keys:
            template_file = os.path.join(self.model_dir, conf_key)
            conf_file = os.path.join(self.work_dir, f'{conf_key}.conf')
            update_conf(template_file, conf_file, custom_conf[conf_key])
            self.conf_files.append(conf_file)
        self.stages = (math.floor(self._max_epochs * EASY_RATIO),
                       math.floor(self._max_epochs * NORMAL_RATIO),
                       math.floor(self._max_epochs * HARD_RATIO))

    def build_model(self) -> nn.Module:
        """ Instantiate a pytorch model and return.

        By default, we will create a model using config from configuration file. You can
        override this method in a subclass.

        """
        model = Model.from_pretrained(
            self.model_dir, cfg_dict=self.cfg, training=True)
        if isinstance(model, TorchModel) and hasattr(model, 'model'):
            return model.model
        elif isinstance(model, nn.Module):
            return model

    def train(self, *args, **kwargs):
        if not self.data_val:
            self.gen_val()
        logger.info('Start training...')
        totaltime = datetime.datetime.now()

        next_stage_head_epoch = 0
        for stage, num_epoch in enumerate(self.stages):
            next_stage_head_epoch += num_epoch
            epochs_to_run = next_stage_head_epoch - self._current_epoch
            self.run_stage(stage, epochs_to_run)

        # total time spent
        totaltime = datetime.datetime.now() - totaltime
        logger.info('Total time spent: {:.2f} hours\n'.format(
            totaltime.total_seconds() / 3600.0))

    def run_stage(self, stage, epochs_to_run):
        """
        Run training stages with correspond data

        Args:
            stage: id of stage
            epochs_to_run: the number of epoch to run in this stage
        """
        if epochs_to_run <= 0:
            logger.warning(f'Invalid epoch number, stage {stage} exit!')
            return
        logger.info(f'Starting stage {stage}...')
        dataset, dataloader = self.create_dataloader(
            self.conf_files[stage * 2], self.conf_files[stage * 2 + 1])
        it = iter(dataloader)
        for _ in range(epochs_to_run):
            self._current_epoch += 1
            epochtime = datetime.datetime.now()
            logger.info('Start epoch %d...', self._current_epoch)
            loss_train_epoch = 0.0
            validbatchs = 0
            for bi in range(self._train_iters):
                # prepare data
                feat, label = next(it)
                label = torch.reshape(label, (-1, ))
                feat = to_device(feat, self.device)
                label = to_device(label, self.device)
                # apply model
                self.optimizer.zero_grad()
                predict = self.model(feat)
                # calculate loss
                loss = self.loss_fn(
                    torch.reshape(predict, (-1, self._num_classes)), label)
                if not np.isnan(loss.item()):
                    loss.backward()
                    self.optimizer.step()
                    loss_train_epoch += loss.item()
                    validbatchs += 1
                train_result = 'Epoch: {:04d}/{:04d}, batch: {:04d}/{:04d}, loss: {:.4f}'.format(
                    self._current_epoch, self._max_epochs, bi + 1,
                    self._train_iters, loss.item())
                logger.info(train_result)
                self._dump_log(train_result)

            # average training loss in one epoch
            loss_train_epoch /= validbatchs
            loss_val_epoch = self.evaluate('')
            val_result = 'Evaluate epoch: {:04d}, loss_train: {:.4f}, loss_val: {:.4f}'.format(
                self._current_epoch, loss_train_epoch, loss_val_epoch)
            logger.info(val_result)
            self._dump_log(val_result)
            # check point
            ckpt_name = '{}_{:04d}_loss_train_{:.4f}_loss_val_{:.4f}.pth'.format(
                CKPT_PREFIX, self._current_epoch, loss_train_epoch,
                loss_val_epoch)
            save_path = os.path.join(self.work_dir, ckpt_name)
            logger.info(f'Save model to {save_path}')
            torch.save(self.model, save_path)
            # time spent per epoch
            epochtime = datetime.datetime.now() - epochtime
            logger.info('Epoch {:04d} time spent: {:.2f} hours'.format(
                self._current_epoch,
                epochtime.total_seconds() / 3600.0))
        dataloader.stop()
        dataset.release()
        logger.info(f'Stage {stage} is finished.')

    def gen_val(self):
        """
        generate validation set
        """
        val_dump_file = os.path.join(self.work_dir, 'val_dataset.bin')
        if self._current_epoch > 0:
            logger.info('Start loading validation set...')
            with open(val_dump_file, 'rb') as f:
                self.data_val = pickle.load(f)
            logger.info('Finish loading validation set!')
            return

        logger.info('Start generating validation set...')
        dataset, dataloader = self.create_dataloader(self.conf_files[2],
                                                     self.conf_files[3])
        it = iter(dataloader)

        self.data_val = []
        for bi in range(self._val_iters):
            logger.info('Iterating validation data %d', bi)
            feat, label = next(it)
            label = torch.reshape(label, (-1, ))
            self.data_val.append([feat, label])

        dataloader.stop()
        dataset.release()

        with open(val_dump_file, 'wb') as f:
            pickle.dump(self.data_val, f)
        logger.info('Finish generating validation set!')

    def create_dataloader(self, base_path, finetune_path):
        dataset = KWSDataset(base_path, finetune_path, self._threads,
                             self._single_rate, self._num_classes)
        dataloader = KWSDataLoader(
            dataset, batchsize=self._batch_size, numworkers=self._threads)
        dataloader.start()
        return dataset, dataloader

    def evaluate(self, checkpoint_path: str, *args,
                 **kwargs) -> Dict[str, float]:
        logger.info('Start validation...')
        loss_val_epoch = 0.0

        with torch.no_grad():
            for feat, label in self.data_val:
                feat = to_device(feat, self.device)
                label = to_device(label, self.device)
                # apply model
                predict = self.model(feat)
                # calculate loss
                loss = self.loss_fn(
                    torch.reshape(predict, (-1, self._num_classes)), label)
                loss_val_epoch += loss.item()
        logger.info('Finish validation.')
        return loss_val_epoch / self._val_iters

    def _dump_log(self, msg):
        if is_master():
            with open(self.json_log_path, 'a+') as f:
                f.write(msg)
                f.write('\n')