mrq/DL-Art-School
1
3
Fork 5
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Add 8x dim reductor

Browse Source
This commit is contained in:
James Betker 2022-05-22 20:23:16 -06:00
parent 07f7be24ce
commit 41809a6330
1 changed files with 23 additions and 8 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

31
codes/models/audio/mel2vec.py
View File

@ -362,16 +362,31 @@ class Mel2Vec(nn.Module):
mask_time_length=10, mask_time_length=10,
disable_custom_linear_init=False, disable_custom_linear_init=False,
linear_init_scale=.02, linear_init_scale=.02,
dim_reduction_multiplier=4,
): ):
super().__init__() super().__init__()
self.input_blocks = nn.Sequential(nn.Conv1d(mel_input_channels, inner_dim//2, kernel_size=5, padding=2, stride=2), if dim_reduction_multiplier == 4:
nn.GroupNorm(num_groups=8, num_channels=inner_dim//2, affine=True), self.input_blocks = nn.Sequential(nn.Conv1d(mel_input_channels, inner_dim//2, kernel_size=5, padding=2, stride=2),
nn.GELU(), nn.GroupNorm(num_groups=8, num_channels=inner_dim//2, affine=True),
nn.Conv1d(inner_dim//2, inner_dim, kernel_size=3, padding=1, stride=2), nn.GELU(),
nn.GELU(), nn.Conv1d(inner_dim//2, inner_dim, kernel_size=3, padding=1, stride=2),
nn.Conv1d(inner_dim, inner_dim, kernel_size=3, padding=1), nn.GELU(),
nn.GELU(), nn.Conv1d(inner_dim, inner_dim, kernel_size=3, padding=1),
) nn.GELU(),
)
elif dim_reduction_multiplier == 8:
self.input_blocks = nn.Sequential(nn.Conv1d(mel_input_channels, inner_dim//4, kernel_size=5, padding=2, stride=2),
nn.GroupNorm(num_groups=8, num_channels=inner_dim//4, affine=True),
nn.GELU(),
nn.Conv1d(inner_dim//4, inner_dim//2, kernel_size=3, padding=1, stride=2),
nn.GELU(),
nn.Conv1d(inner_dim//2, inner_dim, kernel_size=3, padding=1, stride=2),
nn.GELU(),
nn.Conv1d(inner_dim, inner_dim, kernel_size=3, padding=1),
nn.GELU(),
)
else:
assert False, f"dim_reduction_multiplier={dim_reduction_multiplier} not supported"
self.projector = Mel2Vec2FeatureProjection(inner_dim, dropout) self.projector = Mel2Vec2FeatureProjection(inner_dim, dropout)
self.masked_spec_embed = nn.Parameter(torch.rand(inner_dim,)) self.masked_spec_embed = nn.Parameter(torch.rand(inner_dim,))
self.encoder = Wav2Vec2Encoder(inner_dim, dropout, layers, layerdrop) self.encoder = Wav2Vec2Encoder(inner_dim, dropout, layers, layerdrop)