Skip to content

speaker

EncoderClassifier

Bases: Module

Source code in speechain/module/encoder/speaker.py 
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
class EncoderClassifier(nn.Module):
    def __init__(self, model_type="ecapa"):
        super().__init__()
        self.model_type = model_type
        self.model = self._create_model()

    def _create_model(self):
        if self.model_type == "ecapa":
            return self._create_ecapa()
        elif self.model_type == "xvector":
            return self._create_xvector()
        else:
            raise ValueError(f"Unknown model type: {self.model_type}")

    def _create_ecapa(self):
        channels = 512
        model = nn.Sequential(
            nn.Conv1d(80, channels, 7, padding=3),
            nn.BatchNorm1d(channels),
            nn.ReLU(),
            nn.Sequential(
                Res2Block(channels), SEModule(channels), nn.BatchNorm1d(channels)
            ),
            nn.AdaptiveAvgPool1d(1),
            nn.Flatten(),
            nn.Linear(channels, 192),
        )
        return model

    def _create_xvector(self):
        model = nn.Sequential(
            nn.Conv1d(80, 512, 5, padding=2),
            nn.BatchNorm1d(512),
            nn.ReLU(),
            nn.Conv1d(512, 512, 3, padding=1),
            nn.BatchNorm1d(512),
            nn.ReLU(),
            nn.Conv1d(512, 512, 3, padding=1),
            nn.BatchNorm1d(512),
            nn.ReLU(),
            nn.AdaptiveAvgPool1d(1),
            nn.Flatten(),
            nn.Linear(512, 512),
            nn.ReLU(),
            nn.Linear(512, 192),
        )
        return model

    def encode_batch(self, wavs, wav_lens=None):
        self.eval()
        with torch.no_grad():
            x = wavs.transpose(1, 2)
            embeddings = self.model(x)
            return F.normalize(embeddings, p=2, dim=1)

    @classmethod
    def from_hparams(cls, source, savedir=None, run_opts=None):
        """Load pretrained model."""
        model = cls(model_type="ecapa" if "ecapa" in source else "xvector")

        if run_opts and "device" in run_opts:
            model = model.to(run_opts["device"])

        if savedir:
            weights_path = os.path.join(savedir, "encoder.pth")
            if os.path.exists(weights_path):
                model.load_state_dict(
                    torch.load(weights_path, map_location=run_opts["device"])
                )

        return model

from_hparams(source, savedir=None, run_opts=None) classmethod

Load pretrained model.

Source code in speechain/module/encoder/speaker.py 
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
@classmethod
def from_hparams(cls, source, savedir=None, run_opts=None):
    """Load pretrained model."""
    model = cls(model_type="ecapa" if "ecapa" in source else "xvector")

    if run_opts and "device" in run_opts:
        model = model.to(run_opts["device"])

    if savedir:
        weights_path = os.path.join(savedir, "encoder.pth")
        if os.path.exists(weights_path):
            model.load_state_dict(
                torch.load(weights_path, map_location=run_opts["device"])
            )

    return model