From 10378fc37f3f93a84c7fd2b5d0349e8755f62b8f Mon Sep 17 00:00:00 2001 From: James Betker Date: Wed, 18 May 2022 11:07:12 -0600 Subject: [PATCH] make codebooks specifiable --- codes/models/audio/mel2vec.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/codes/models/audio/mel2vec.py b/codes/models/audio/mel2vec.py index ff903da8..e02768af 100644 --- a/codes/models/audio/mel2vec.py +++ b/codes/models/audio/mel2vec.py @@ -541,7 +541,9 @@ class Wav2Vec2GumbelVectorQuantizer(nn.Module): class ContrastiveTrainingWrapper(nn.Module): def __init__(self, inner_dim=1024, dropout=.1, mask_time_prob=.5, mask_time_length=6, num_negatives=100, - max_gumbel_temperature=2.0, min_gumbel_temperature=.5, gumbel_temperature_decay=.999995, **kwargs): + max_gumbel_temperature=2.0, min_gumbel_temperature=.5, gumbel_temperature_decay=.999995, + codebook_size=320, codebook_groups=2, + **kwargs): super().__init__() self.m2v = Mel2Vec(inner_dim=inner_dim, dropout=dropout, mask_time_prob=mask_time_prob, mask_time_length=mask_time_length, **kwargs) @@ -551,7 +553,7 @@ class ContrastiveTrainingWrapper(nn.Module): self.max_gumbel_temperature = max_gumbel_temperature self.min_gumbel_temperature = min_gumbel_temperature self.gumbel_temperature_decay = gumbel_temperature_decay - self.quantizer = Wav2Vec2GumbelVectorQuantizer(inner_dim) + self.quantizer = Wav2Vec2GumbelVectorQuantizer(inner_dim, num_codevector_groups=codebook_groups, num_codevectors_per_group=codebook_size) self.num_losses_record = [] # make sure that project_hid & project_q are initialized like normal linear layers