Method report_core

funasr/models/eend/utils/report.py:149–202

Report core. Args: model: Model instance or model name. data_loader: TODO. device: Target device ("cuda:0", "cpu", etc.).

(self, model, data_loader, device)

        return num

    def report_core(self, model, data_loader, device):
        """Report core.
        
            Args:
                model: Model instance or model name.
                data_loader: TODO.
                device: Target device ("cuda:0", "cpu", etc.).
            """
        res = {}
        for item in metrics:
            res[item[0]] = 0.0
            res[item[1]] = 0.0
        with torch.no_grad():
            loss_s = 0.0
            uidx = 0
            for xs, ts, orders in data_loader:
                xs = [x.to(device) for x in xs]
                ts = [t.to(device) for t in ts]
                orders = [o.to(device) for o in orders]
                loss, pit_loss, mpit_loss, att_loss, ys, logits, labels, attractors = model(
                    xs, ts, orders
                )
                loss_s += loss.item()
                uidx += 1

                for logit, t, att in zip(logits, labels, attractors):
                    pred = torch.argmax(torch.softmax(logit, dim=-1), dim=-1)  # (T, )
                    oov_index = torch.where(pred == self.mapping_dict["oov"])[0]
                    for i in oov_index:
                        if i > 0:
                            pred[i] = pred[i - 1]
                        else:
                            pred[i] = 0
                    pred = [self.inv_mapping_func(i, self.mapping_dict) for i in pred]
                    decisions = [bin(num)[2:].zfill(self.max_n_speaker)[::-1] for num in pred]
                    decisions = (
                        torch.from_numpy(
                            np.stack([np.array([int(i) for i in dec]) for dec in decisions], axis=0)
                        )
                        .to(att.device)
                        .to(torch.float32)
                    )
                    decisions = decisions[:, : att.shape[0]]

                    stats = self.calc_diarization_error(decisions, t)
                    res["speaker_scored"] += stats["speaker_scored"]
                    res["speech_scored"] += stats["speech_scored"]
                    res["frames"] += stats["frames"]
                    for item in metrics:
                        res[item[0]] += stats[item[0]]
                loss_s /= uidx
                vad_acc = 0

        return res, loss_s, stats.keys(), vad_acc

    def calc_diarization_error(self, decisions, label, label_delay=0):
        """Calc diarization error.