# -*- coding: utf-8 -*- import os import os.path as osp import sys import time from collections import defaultdict import numpy as np import torch from torch import nn from PIL import Image from tqdm import tqdm from utils import calc_wer import logging logger = logging.getLogger(__name__) logger.setLevel(logging.DEBUG) from utils import * class Trainer(object): def __init__(self, model=None, criterion=None, optimizer=None, scheduler=None, config={}, device=torch.device("cpu"), logger=logger, train_dataloader=None, val_dataloader=None, initial_steps=0, initial_epochs=0): self.steps = initial_steps self.epochs = initial_epochs self.model = model self.criterion = criterion self.optimizer = optimizer self.scheduler = scheduler self.train_dataloader = train_dataloader self.val_dataloader = val_dataloader self.config = config self.device = device self.finish_train = False self.logger = logger self.fp16_run = False def save_checkpoint(self, checkpoint_path): """Save checkpoint. Args: checkpoint_path (str): Checkpoint path to be saved. """ state_dict = { "optimizer": self.optimizer.state_dict(), "scheduler": self.scheduler.state_dict(), "steps": self.steps, "epochs": self.epochs, } state_dict["model"] = self.model.state_dict() if not os.path.exists(os.path.dirname(checkpoint_path)): os.makedirs(os.path.dirname(checkpoint_path)) torch.save(state_dict, checkpoint_path) def load_checkpoint(self, checkpoint_path, load_only_params=False): """Load checkpoint. Args: checkpoint_path (str): Checkpoint path to be loaded. load_only_params (bool): Whether to load only model parameters. """ state_dict = torch.load(checkpoint_path, map_location="cpu") self._load(state_dict["model"], self.model) if not load_only_params: self.steps = state_dict["steps"] self.epochs = state_dict["epochs"] self.optimizer.load_state_dict(state_dict["optimizer"]) # overwrite schedular argument parameters state_dict["scheduler"].update(**self.config.get("scheduler_params", {})) self.scheduler.load_state_dict(state_dict["scheduler"]) def _load(self, states, model, force_load=True): model_states = model.state_dict() for key, val in states.items(): try: if key not in model_states: continue if isinstance(val, nn.Parameter): val = val.data if val.shape != model_states[key].shape: self.logger.info("%s does not have same shape" % key) print(val.shape, model_states[key].shape) if not force_load: continue min_shape = np.minimum(np.array(val.shape), np.array(model_states[key].shape)) slices = [slice(0, min_index) for min_index in min_shape] model_states[key][slices].copy_(val[slices]) else: model_states[key].copy_(val) except: self.logger.info("not exist :%s" % key) print("not exist ", key) @staticmethod def get_gradient_norm(model): total_norm = 0 for p in model.parameters(): param_norm = p.grad.data.norm(2) total_norm += param_norm.item() ** 2 total_norm = np.sqrt(total_norm) return total_norm @staticmethod def length_to_mask(lengths): mask = torch.arange(lengths.max()).unsqueeze(0).expand(lengths.shape[0], -1).type_as(lengths) mask = torch.gt(mask+1, lengths.unsqueeze(1)) return mask def _get_lr(self): for param_group in self.optimizer.param_groups: lr = param_group['lr'] break return lr @staticmethod def get_image(arrs): pil_images = [] height = 0 width = 0 for arr in arrs: uint_arr = (((arr - arr.min()) / (arr.max() - arr.min())) * 255).astype(np.uint8) pil_image = Image.fromarray(uint_arr) pil_images.append(pil_image) height += uint_arr.shape[0] width = max(width, uint_arr.shape[1]) palette = Image.new('L', (width, height)) curr_heigth = 0 for pil_image in pil_images: palette.paste(pil_image, (0, curr_heigth)) curr_heigth += pil_image.size[1] return palette def run(self, batch): self.optimizer.zero_grad() batch = [b.to(self.device) for b in batch] text_input, text_input_length, mel_input, mel_input_length = batch mel_input_length = mel_input_length // (2 ** self.model.n_down) future_mask = self.model.get_future_mask( mel_input.size(2)//(2**self.model.n_down), unmask_future_steps=0).to(self.device) mel_mask = self.model.length_to_mask(mel_input_length) text_mask = self.model.length_to_mask(text_input_length) ppgs, s2s_pred, s2s_attn = self.model( mel_input, src_key_padding_mask=mel_mask, text_input=text_input) loss_ctc = self.criterion['ctc'](ppgs.log_softmax(dim=2).transpose(0, 1), text_input, mel_input_length, text_input_length) loss_s2s = 0 for _s2s_pred, _text_input, _text_length in zip(s2s_pred, text_input, text_input_length): loss_s2s += self.criterion['ce'](_s2s_pred[:_text_length], _text_input[:_text_length]) loss_s2s /= text_input.size(0) loss = loss_ctc + loss_s2s loss.backward() torch.nn.utils.clip_grad_value_(self.model.parameters(), 5) self.optimizer.step() self.scheduler.step() return {'loss': loss.item(), 'ctc': loss_ctc.item(), 's2s': loss_s2s.item()} def _train_epoch(self): train_losses = defaultdict(list) self.model.train() for train_steps_per_epoch, batch in enumerate(tqdm(self.train_dataloader, desc="[train]"), 1): losses = self.run(batch) for key, value in losses.items(): train_losses["train/%s" % key].append(value) train_losses = {key: np.mean(value) for key, value in train_losses.items()} train_losses['train/learning_rate'] = self._get_lr() return train_losses @torch.no_grad() def _eval_epoch(self): self.model.eval() eval_losses = defaultdict(list) eval_images = defaultdict(list) for eval_steps_per_epoch, batch in enumerate(tqdm(self.val_dataloader, desc="[eval]"), 1): batch = [b.to(self.device) for b in batch] text_input, text_input_length, mel_input, mel_input_length = batch mel_input_length = mel_input_length // (2 ** self.model.n_down) future_mask = self.model.get_future_mask( mel_input.size(2)//(2**self.model.n_down), unmask_future_steps=0).to(self.device) mel_mask = self.model.length_to_mask(mel_input_length) text_mask = self.model.length_to_mask(text_input_length) ppgs, s2s_pred, s2s_attn = self.model( mel_input, src_key_padding_mask=mel_mask, text_input=text_input) loss_ctc = self.criterion['ctc'](ppgs.log_softmax(dim=2).transpose(0, 1), text_input, mel_input_length, text_input_length) loss_s2s = 0 for _s2s_pred, _text_input, _text_length in zip(s2s_pred, text_input, text_input_length): loss_s2s += self.criterion['ce'](_s2s_pred[:_text_length], _text_input[:_text_length]) loss_s2s /= text_input.size(0) loss = loss_ctc + loss_s2s eval_losses["eval/ctc"].append(loss_ctc.item()) eval_losses["eval/s2s"].append(loss_s2s.item()) eval_losses["eval/loss"].append(loss.item()) _, amax_ppgs = torch.max(ppgs, dim=2) wers = [calc_wer(target[:text_length], pred[:mel_length], ignore_indexes=list(range(5))) \ for target, pred, text_length, mel_length in zip( text_input.cpu(), amax_ppgs.cpu(), text_input_length.cpu(), mel_input_length.cpu())] eval_losses["eval/wer"].extend(wers) _, amax_s2s = torch.max(s2s_pred, dim=2) acc = [torch.eq(target[:length], pred[:length]).float().mean().item() \ for target, pred, length in zip(text_input.cpu(), amax_s2s.cpu(), text_input_length.cpu())] eval_losses["eval/acc"].extend(acc) if eval_steps_per_epoch <= 2: eval_images["eval/image"].append( self.get_image([s2s_attn[0].cpu().numpy()])) eval_losses = {key: np.mean(value) for key, value in eval_losses.items()} eval_losses.update(eval_images) return eval_losses