import torch import torchaudio from cosine import WarmupCosineScheduler from datamodule.transforms import TextTransform from espnet.nets.batch_beam_search import BatchBeamSearch from espnet.nets.pytorch_backend.e2e_asr_conformer import E2E from espnet.nets.scorers.length_bonus import LengthBonus from pytorch_lightning import LightningModule def compute_word_level_distance(seq1, seq2): seq1, seq2 = seq1.lower().split(), seq2.lower().split() return torchaudio.functional.edit_distance(seq1, seq2) class ModelModule(LightningModule): def __init__(self, args): super().__init__() self.args = args self.save_hyperparameters(args) self.modality = args.modality self.text_transform = TextTransform() self.token_list = self.text_transform.token_list self.model = E2E(len(self.token_list), self.modality, ctc_weight=getattr(args, "ctc_weight", 0.1)) # -- initialise if getattr(args, "pretrained_model_path", None): ckpt = torch.load(args.pretrained_model_path, map_location=lambda storage, loc: storage) if getattr(args, "transfer_frontend", False): tmp_ckpt = {k: v for k, v in ckpt["model_state_dict"].items() if k.startswith("trunk.") or k.startswith("frontend3D.")} self.model.frontend.load_state_dict(tmp_ckpt) print("Pretrained weights of the frontend component are loaded successfully.") elif getattr(args, "transfer_encoder", False): tmp_ckpt = {k.replace("frontend.",""):v for k,v in ckpt.items() if k.startswith("frontend.")} self.model.frontend.load_state_dict(tmp_ckpt) tmp_ckpt = {k.replace("proj_encoder.",""):v for k,v in ckpt.items() if k.startswith("proj_encoder.")} self.model.proj_encoder.load_state_dict(tmp_ckpt) tmp_ckpt = {k.replace("encoder.",""):v for k,v in ckpt.items() if k.startswith("encoder.")} self.model.encoder.load_state_dict(tmp_ckpt) print("Pretrained weights of the frontend, proj_encoder and encoder component are loaded successfully.") else: self.model.load_state_dict(ckpt) print("Pretrained weights of the full model are loaded successfully.") def configure_optimizers(self): optimizer = torch.optim.AdamW(self.model.parameters(), lr=self.args.lr, weight_decay=self.args.weight_decay, betas=(0.9, 0.98)) scheduler = WarmupCosineScheduler(optimizer, self.args.warmup_epochs, self.args.max_epochs, len(self.trainer.datamodule.train_dataloader()) / self.trainer.num_devices / self.trainer.num_nodes) scheduler = {"scheduler": scheduler, "interval": "step"} return [optimizer], [scheduler] def forward(self, sample): self.beam_search = get_beam_search_decoder(self.model, self.token_list) x = self.model.frontend(sample.unsqueeze(0)) x = self.model.proj_encoder(x) enc_feat, _ = self.model.encoder(x, None) enc_feat = enc_feat.squeeze(0) nbest_hyps = self.beam_search(enc_feat) nbest_hyps = [h.asdict() for h in nbest_hyps[: min(len(nbest_hyps), 1)]] predicted_token_id = torch.tensor(list(map(int, nbest_hyps[0]["yseq"][1:]))) predicted = self.text_transform.post_process(predicted_token_id).replace("", "") return predicted def validation_step(self, batch, batch_idx): return self._step(batch, batch_idx, step_type="val") def test_step(self, sample, sample_idx): x = self.model.frontend(sample["input"].unsqueeze(0)) x = self.model.proj_encoder(x) enc_feat, _ = self.model.encoder(x, None) enc_feat = enc_feat.squeeze(0) nbest_hyps = self.beam_search(enc_feat) nbest_hyps = [h.asdict() for h in nbest_hyps[: min(len(nbest_hyps), 1)]] predicted_token_id = torch.tensor(list(map(int, nbest_hyps[0]["yseq"][1:]))) predicted = self.text_transform.post_process(predicted_token_id).replace("", "") actual_token_id = sample["target"] actual = self.text_transform.post_process(actual_token_id) self.total_edit_distance += compute_word_level_distance(actual, predicted) self.total_length += len(actual.split()) return def training_step(self, batch, batch_idx): loss = self._step(batch, batch_idx, "train") batch_size = batch["inputs"].size(0) batch_sizes = self.all_gather(batch_size) loss *= batch_sizes.size(0) / batch_sizes.sum() # world size / batch size self.log("monitoring_step", torch.tensor(self.global_step, dtype=torch.float32)) return loss def _step(self, batch, batch_idx, step_type): loss, loss_ctc, loss_att, acc = self.model(batch["inputs"], batch["input_lengths"], batch["targets"]) batch_size = len(batch["inputs"]) if step_type == "train": self.log("loss", loss, on_step=True, on_epoch=True, batch_size=batch_size) self.log("loss_ctc", loss_ctc, on_step=False, on_epoch=True, batch_size=batch_size, sync_dist=True) self.log("loss_att", loss_att, on_step=False, on_epoch=True, batch_size=batch_size, sync_dist=True) self.log("decoder_acc", acc, on_step=True, on_epoch=True, batch_size=batch_size, sync_dist=True) else: self.log("loss_val", loss, batch_size=batch_size, sync_dist=True) self.log("loss_ctc_val", loss_ctc, batch_size=batch_size, sync_dist=True) self.log("loss_att_val", loss_att, batch_size=batch_size, sync_dist=True) self.log("decoder_acc_val", acc, batch_size=batch_size, sync_dist=True) if step_type == "train": self.log("monitoring_step", torch.tensor(self.global_step, dtype=torch.float32)) return loss def on_test_epoch_start(self): self.total_length = 0 self.total_edit_distance = 0 self.text_transform = TextTransform() self.beam_search = get_beam_search_decoder(self.model, self.token_list) def on_test_epoch_end(self): self.log("wer", self.total_edit_distance / self.total_length) def get_beam_search_decoder( model, token_list, rnnlm=None, rnnlm_conf=None, penalty=0, ctc_weight=0.1, lm_weight=0.0, beam_size=40, ): sos = model.odim - 1 eos = model.odim - 1 scorers = model.scorers() scorers["lm"] = None scorers["length_bonus"] = LengthBonus(len(token_list)) weights = { "decoder": 1.0 - ctc_weight, "ctc": ctc_weight, "lm": lm_weight, "length_bonus": penalty, } return BatchBeamSearch( beam_size=beam_size, vocab_size=len(token_list), weights=weights, scorers=scorers, sos=sos, eos=eos, token_list=token_list, pre_beam_score_key=None if ctc_weight == 1.0 else "decoder", )