maai 0.0.12__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- maai-0.0.12/.gitignore +12 -0
- maai-0.0.12/LICENSE +21 -0
- maai-0.0.12/PKG-INFO +40 -0
- maai-0.0.12/maai/__init__.py +4 -0
- maai-0.0.12/maai/encoder.py +88 -0
- maai-0.0.12/maai/encoder_components.py +511 -0
- maai-0.0.12/maai/input.py +240 -0
- maai-0.0.12/maai/model.py +268 -0
- maai-0.0.12/maai/models/config.py +56 -0
- maai-0.0.12/maai/models/vap.py +100 -0
- maai-0.0.12/maai/models/vap_bc.py +99 -0
- maai-0.0.12/maai/models/vap_nod.py +99 -0
- maai-0.0.12/maai/modules.py +501 -0
- maai-0.0.12/maai/objective.py +625 -0
- maai-0.0.12/maai/output.py +242 -0
- maai-0.0.12/maai/util.py +395 -0
- maai-0.0.12/pyproject.toml +42 -0
maai-0.0.12/.gitignore
ADDED
maai-0.0.12/LICENSE
ADDED
|
@@ -0,0 +1,21 @@
|
|
|
1
|
+
MIT License
|
|
2
|
+
|
|
3
|
+
Copyright (c) 2024 Koji Inoue
|
|
4
|
+
|
|
5
|
+
Permission is hereby granted, free of charge, to any person obtaining a copy
|
|
6
|
+
of this software and associated documentation files (the "Software"), to deal
|
|
7
|
+
in the Software without restriction, including without limitation the rights
|
|
8
|
+
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
|
9
|
+
copies of the Software, and to permit persons to whom the Software is
|
|
10
|
+
furnished to do so, subject to the following conditions:
|
|
11
|
+
|
|
12
|
+
The above copyright notice and this permission notice shall be included in all
|
|
13
|
+
copies or substantial portions of the Software.
|
|
14
|
+
|
|
15
|
+
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
|
16
|
+
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
|
17
|
+
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
|
18
|
+
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
|
19
|
+
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
|
20
|
+
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
|
21
|
+
SOFTWARE.
|
maai-0.0.12/PKG-INFO
ADDED
|
@@ -0,0 +1,40 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: maai
|
|
3
|
+
Version: 0.0.12
|
|
4
|
+
Summary: Real-time and Continuous Non-Linguistic Behavior (Maai) Generation Software
|
|
5
|
+
Project-URL: Homepage, https://github.com/MaAI-Kyoto/MaAI
|
|
6
|
+
Author-email: MaAI team <inoue@sap.ist.i.kyoto-u.ac.jp>
|
|
7
|
+
License: MIT License
|
|
8
|
+
|
|
9
|
+
Copyright (c) 2024 Koji Inoue
|
|
10
|
+
|
|
11
|
+
Permission is hereby granted, free of charge, to any person obtaining a copy
|
|
12
|
+
of this software and associated documentation files (the "Software"), to deal
|
|
13
|
+
in the Software without restriction, including without limitation the rights
|
|
14
|
+
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
|
15
|
+
copies of the Software, and to permit persons to whom the Software is
|
|
16
|
+
furnished to do so, subject to the following conditions:
|
|
17
|
+
|
|
18
|
+
The above copyright notice and this permission notice shall be included in all
|
|
19
|
+
copies or substantial portions of the Software.
|
|
20
|
+
|
|
21
|
+
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
|
22
|
+
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
|
23
|
+
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
|
24
|
+
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
|
25
|
+
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
|
26
|
+
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
|
27
|
+
SOFTWARE.
|
|
28
|
+
License-File: LICENSE
|
|
29
|
+
Requires-Python: >=3.10
|
|
30
|
+
Requires-Dist: einops==0.7.0
|
|
31
|
+
Requires-Dist: fastapi==0.111.0
|
|
32
|
+
Requires-Dist: huggingface-hub
|
|
33
|
+
Requires-Dist: matplotlib==3.7.5
|
|
34
|
+
Requires-Dist: numpy
|
|
35
|
+
Requires-Dist: pyaudio
|
|
36
|
+
Requires-Dist: pydub==0.25.1
|
|
37
|
+
Requires-Dist: pygame
|
|
38
|
+
Requires-Dist: seaborn==0.13.2
|
|
39
|
+
Requires-Dist: soundfile==0.12.1
|
|
40
|
+
Requires-Dist: torch>=2.2.0
|
|
@@ -0,0 +1,88 @@
|
|
|
1
|
+
import torch
|
|
2
|
+
import torch.nn as nn
|
|
3
|
+
import einops
|
|
4
|
+
import os
|
|
5
|
+
import numpy as np
|
|
6
|
+
|
|
7
|
+
from .encoder_components import load_CPC, get_cnn_layer
|
|
8
|
+
|
|
9
|
+
import time
|
|
10
|
+
|
|
11
|
+
class EncoderCPC(nn.Module):
|
|
12
|
+
"""
|
|
13
|
+
Encoder: waveform -> h
|
|
14
|
+
pretrained: default='cpc'
|
|
15
|
+
|
|
16
|
+
A simpler version of the Encoder
|
|
17
|
+
check paper (branch) version to see other encoders...
|
|
18
|
+
"""
|
|
19
|
+
|
|
20
|
+
def __init__(self, load_pretrained=True, freeze=True, cpc_model=''):
|
|
21
|
+
|
|
22
|
+
super().__init__()
|
|
23
|
+
|
|
24
|
+
self.sample_rate = 16000
|
|
25
|
+
|
|
26
|
+
if load_pretrained:
|
|
27
|
+
self.encoder = load_CPC(checkpoint_cpc=cpc_model, load_state_dict=True)
|
|
28
|
+
else:
|
|
29
|
+
self.encoder = load_CPC(checkpoint_cpc='', load_state_dict=False)
|
|
30
|
+
|
|
31
|
+
# Keep Hidden layer
|
|
32
|
+
self.encoder.gAR.keepHidden = True
|
|
33
|
+
|
|
34
|
+
self.output_dim = self.encoder.gEncoder.conv4.out_channels
|
|
35
|
+
self.dim = self.output_dim
|
|
36
|
+
|
|
37
|
+
self.downsample_ratio = 160
|
|
38
|
+
self.downsample = get_cnn_layer(
|
|
39
|
+
dim=self.output_dim,
|
|
40
|
+
kernel=[5],
|
|
41
|
+
stride=[2],
|
|
42
|
+
dilation=[1],
|
|
43
|
+
activation="GELU",
|
|
44
|
+
)
|
|
45
|
+
self.downsample_ratio = 320
|
|
46
|
+
|
|
47
|
+
if freeze:
|
|
48
|
+
self.freeze()
|
|
49
|
+
|
|
50
|
+
def get_default_conf(self):
|
|
51
|
+
return {""}
|
|
52
|
+
|
|
53
|
+
def freeze(self):
|
|
54
|
+
for p in self.encoder.parameters():
|
|
55
|
+
p.requires_grad_(False)
|
|
56
|
+
print(f"Froze {self.__class__.__name__}!")
|
|
57
|
+
|
|
58
|
+
def unfreeze(self):
|
|
59
|
+
for p in self.encoder.parameters():
|
|
60
|
+
p.requires_grad_(True)
|
|
61
|
+
print(f"Trainable {self.__class__.__name__}!")
|
|
62
|
+
|
|
63
|
+
def forward(self, waveform):
|
|
64
|
+
|
|
65
|
+
if waveform.ndim < 3:
|
|
66
|
+
waveform = waveform.unsqueeze(1) # channel dim
|
|
67
|
+
|
|
68
|
+
# Backwards using only the encoder encounters:
|
|
69
|
+
# ---------------------------------------------------
|
|
70
|
+
# RuntimeError: one of the variables needed for gradient computation
|
|
71
|
+
# has been modified by an inplace operation:
|
|
72
|
+
# [torch.FloatTensor [4, 256, 1000]], which is output 0 of ReluBackward0, is at version 1;
|
|
73
|
+
# expected version 0 instead. Hint: enable anomaly detection to find
|
|
74
|
+
# the operation that failed to compute its gradient, with
|
|
75
|
+
# torch.autograd.set_detect_anomaly(True).
|
|
76
|
+
# HOWEVER, if we feed through encoder.gAR we do not encounter that problem...
|
|
77
|
+
|
|
78
|
+
|
|
79
|
+
z = self.encoder.gEncoder(waveform)
|
|
80
|
+
z = einops.rearrange(z, "b c n -> b n c")
|
|
81
|
+
z = z[:, 1:-1, :]
|
|
82
|
+
z = self.encoder.gAR(z)
|
|
83
|
+
z = self.downsample(z)
|
|
84
|
+
|
|
85
|
+
return z
|
|
86
|
+
|
|
87
|
+
def hash_tensor(self, tensor):
|
|
88
|
+
return hash(tuple(tensor.reshape(-1).tolist()))
|
|
@@ -0,0 +1,511 @@
|
|
|
1
|
+
import argparse
|
|
2
|
+
import torch
|
|
3
|
+
import torch.nn as nn
|
|
4
|
+
import torch.nn.functional as F
|
|
5
|
+
from einops.layers.torch import Rearrange
|
|
6
|
+
from os.path import exists, dirname
|
|
7
|
+
from os import makedirs
|
|
8
|
+
from typing import List
|
|
9
|
+
|
|
10
|
+
#from vap.utils import repo_root
|
|
11
|
+
|
|
12
|
+
"""
|
|
13
|
+
|
|
14
|
+
#############################################################
|
|
15
|
+
#############################################################
|
|
16
|
+
WARNING - ATTENTION - HEAR YEE HEAAAR YEEEE
|
|
17
|
+
#############################################################
|
|
18
|
+
#############################################################
|
|
19
|
+
|
|
20
|
+
Most of the code in this file are scaled down (and heavily copied) versions of
|
|
21
|
+
|
|
22
|
+
https://github.com/facebookresearch/CPC_audio
|
|
23
|
+
|
|
24
|
+
Please checkout their codebase if you are interested in CPC networks
|
|
25
|
+
----------------------------------------------------------------------------
|
|
26
|
+
|
|
27
|
+
torch.hub downloads to: `$HOME/.cache/torch/hub/checkpoints/`
|
|
28
|
+
Explicit checkpoint path saved manually in "assets/" see CHECKPOINTS below.
|
|
29
|
+
"""
|
|
30
|
+
|
|
31
|
+
# CHECKPOINTS = {
|
|
32
|
+
# "cpc": "../asset/cpc/60k_epoch4-d0f474de.pt"
|
|
33
|
+
# }
|
|
34
|
+
# NAMES = list(CHECKPOINTS.keys())
|
|
35
|
+
|
|
36
|
+
|
|
37
|
+
class ChannelNorm(nn.Module):
|
|
38
|
+
"""
|
|
39
|
+
Most of the code in this file are scaled down (and heavily copied) versions of
|
|
40
|
+
https://github.com/facebookresearch/CPC_audio
|
|
41
|
+
"""
|
|
42
|
+
|
|
43
|
+
def __init__(self, numFeatures, epsilon=1e-05, affine=True):
|
|
44
|
+
|
|
45
|
+
super(ChannelNorm, self).__init__()
|
|
46
|
+
if affine:
|
|
47
|
+
self.weight = nn.parameter.Parameter(torch.Tensor(1, numFeatures, 1))
|
|
48
|
+
self.bias = nn.parameter.Parameter(torch.Tensor(1, numFeatures, 1))
|
|
49
|
+
else:
|
|
50
|
+
self.weight = None
|
|
51
|
+
self.bias = None
|
|
52
|
+
self.epsilon = epsilon
|
|
53
|
+
self.p = 0
|
|
54
|
+
self.affine = affine
|
|
55
|
+
self.reset_parameters()
|
|
56
|
+
|
|
57
|
+
def reset_parameters(self):
|
|
58
|
+
if self.affine:
|
|
59
|
+
torch.nn.init.ones_(self.weight)
|
|
60
|
+
torch.nn.init.zeros_(self.bias)
|
|
61
|
+
|
|
62
|
+
def forward(self, x):
|
|
63
|
+
|
|
64
|
+
cumMean = x.mean(dim=1, keepdim=True)
|
|
65
|
+
cumVar = x.var(dim=1, keepdim=True)
|
|
66
|
+
x = (x - cumMean) * torch.rsqrt(cumVar + self.epsilon)
|
|
67
|
+
|
|
68
|
+
if self.weight is not None:
|
|
69
|
+
x = x * self.weight + self.bias
|
|
70
|
+
return x
|
|
71
|
+
|
|
72
|
+
|
|
73
|
+
class CPCEncoder(nn.Module):
|
|
74
|
+
"""
|
|
75
|
+
Most of the code in this file are scaled down (and heavily copied) versions of
|
|
76
|
+
https://github.com/facebookresearch/CPC_audio
|
|
77
|
+
"""
|
|
78
|
+
|
|
79
|
+
def __init__(self, sizeHidden=512, normMode="layerNorm"):
|
|
80
|
+
super(CPCEncoder, self).__init__()
|
|
81
|
+
normLayer = ChannelNorm
|
|
82
|
+
self.dimEncoded = sizeHidden
|
|
83
|
+
self.conv0 = nn.Conv1d(1, sizeHidden, 10, stride=5, padding=3)
|
|
84
|
+
self.batchNorm0 = normLayer(sizeHidden)
|
|
85
|
+
self.conv1 = nn.Conv1d(sizeHidden, sizeHidden, 8, stride=4, padding=2)
|
|
86
|
+
self.batchNorm1 = normLayer(sizeHidden)
|
|
87
|
+
self.conv2 = nn.Conv1d(sizeHidden, sizeHidden, 4, stride=2, padding=1)
|
|
88
|
+
self.batchNorm2 = normLayer(sizeHidden)
|
|
89
|
+
self.conv3 = nn.Conv1d(sizeHidden, sizeHidden, 4, stride=2, padding=1)
|
|
90
|
+
self.batchNorm3 = normLayer(sizeHidden)
|
|
91
|
+
self.conv4 = nn.Conv1d(sizeHidden, sizeHidden, 4, stride=2, padding=1)
|
|
92
|
+
self.batchNorm4 = normLayer(sizeHidden)
|
|
93
|
+
self.DOWNSAMPLING = 160
|
|
94
|
+
|
|
95
|
+
def getDimOutput(self):
|
|
96
|
+
return self.conv4.out_channels
|
|
97
|
+
|
|
98
|
+
def forward(self, x):
|
|
99
|
+
x = F.relu(self.batchNorm0(self.conv0(x)))
|
|
100
|
+
x = F.relu(self.batchNorm1(self.conv1(x)))
|
|
101
|
+
x = F.relu(self.batchNorm2(self.conv2(x)))
|
|
102
|
+
x = F.relu(self.batchNorm3(self.conv3(x)))
|
|
103
|
+
x = F.relu(self.batchNorm4(self.conv4(x)))
|
|
104
|
+
return x
|
|
105
|
+
|
|
106
|
+
|
|
107
|
+
class CPCAR(nn.Module):
|
|
108
|
+
"""
|
|
109
|
+
Most of the code in this file are scaled down (and heavily copied) versions of
|
|
110
|
+
https://github.com/facebookresearch/CPC_audio
|
|
111
|
+
"""
|
|
112
|
+
|
|
113
|
+
def __init__(
|
|
114
|
+
self, dimEncoded, dimOutput, keepHidden, nLevelsGRU, mode="GRU", reverse=False
|
|
115
|
+
):
|
|
116
|
+
|
|
117
|
+
super(CPCAR, self).__init__()
|
|
118
|
+
self.RESIDUAL_STD = 0.1
|
|
119
|
+
|
|
120
|
+
if mode == "LSTM":
|
|
121
|
+
self.baseNet = nn.LSTM(
|
|
122
|
+
dimEncoded, dimOutput, num_layers=nLevelsGRU, batch_first=True
|
|
123
|
+
)
|
|
124
|
+
elif mode == "RNN":
|
|
125
|
+
self.baseNet = nn.RNN(
|
|
126
|
+
dimEncoded, dimOutput, num_layers=nLevelsGRU, batch_first=True
|
|
127
|
+
)
|
|
128
|
+
else:
|
|
129
|
+
self.baseNet = nn.GRU(
|
|
130
|
+
dimEncoded, dimOutput, num_layers=nLevelsGRU, batch_first=True
|
|
131
|
+
)
|
|
132
|
+
|
|
133
|
+
self.hidden = None
|
|
134
|
+
self.keepHidden = keepHidden
|
|
135
|
+
self.reverse = reverse
|
|
136
|
+
|
|
137
|
+
def getDimOutput(self):
|
|
138
|
+
return self.baseNet.hidden_size
|
|
139
|
+
|
|
140
|
+
def forward(self, x):
|
|
141
|
+
|
|
142
|
+
if self.reverse:
|
|
143
|
+
x = torch.flip(x, [1])
|
|
144
|
+
try:
|
|
145
|
+
self.baseNet.flatten_parameters()
|
|
146
|
+
except RuntimeError:
|
|
147
|
+
pass
|
|
148
|
+
x, h = self.baseNet(x, self.hidden)
|
|
149
|
+
if self.keepHidden:
|
|
150
|
+
if isinstance(h, tuple):
|
|
151
|
+
self.hidden = tuple(x.detach() for x in h)
|
|
152
|
+
else:
|
|
153
|
+
self.hidden = h.detach()
|
|
154
|
+
|
|
155
|
+
# For better modularity, a sequence's order should be preserved
|
|
156
|
+
# by each module
|
|
157
|
+
if self.reverse:
|
|
158
|
+
x = torch.flip(x, [1])
|
|
159
|
+
return x
|
|
160
|
+
|
|
161
|
+
|
|
162
|
+
class CPCModel(nn.Module):
|
|
163
|
+
"""
|
|
164
|
+
Most of the code in this file are scaled down (and heavily copied) versions of
|
|
165
|
+
https://github.com/facebookresearch/CPC_audio
|
|
166
|
+
"""
|
|
167
|
+
|
|
168
|
+
def __init__(self, encoder, AR):
|
|
169
|
+
super(CPCModel, self).__init__()
|
|
170
|
+
self.gEncoder = encoder
|
|
171
|
+
self.gAR = AR
|
|
172
|
+
|
|
173
|
+
def forward(self, batchData, label):
|
|
174
|
+
encodedData = self.gEncoder(batchData).permute(0, 2, 1)
|
|
175
|
+
cFeature = self.gAR(encodedData)
|
|
176
|
+
return cFeature, encodedData, label
|
|
177
|
+
|
|
178
|
+
|
|
179
|
+
def load_CPC(checkpoint_cpc, load_state_dict=True):
|
|
180
|
+
"""
|
|
181
|
+
Contrast predictive learning model for audio data
|
|
182
|
+
pretrained: if True, load a model trained on libri-light 60k
|
|
183
|
+
(https://arxiv.org/abs/1912.07875)
|
|
184
|
+
**kwargs : see cpc/cpc_default_config to get the list of possible arguments
|
|
185
|
+
|
|
186
|
+
Most of the code in this file are scaled down (and heavily copied) versions of
|
|
187
|
+
https://github.com/facebookresearch/CPC_audio
|
|
188
|
+
"""
|
|
189
|
+
|
|
190
|
+
def loadArgs(args, locArgs, forbiddenAttr=None):
|
|
191
|
+
for k, v in vars(locArgs).items():
|
|
192
|
+
if forbiddenAttr is not None:
|
|
193
|
+
if k not in forbiddenAttr:
|
|
194
|
+
setattr(args, k, v)
|
|
195
|
+
else:
|
|
196
|
+
setattr(args, k, v)
|
|
197
|
+
|
|
198
|
+
def get_default_cpc_config():
|
|
199
|
+
parser = argparse.ArgumentParser()
|
|
200
|
+
|
|
201
|
+
# Run parameters
|
|
202
|
+
group = parser.add_argument_group(
|
|
203
|
+
"Architecture configuration",
|
|
204
|
+
description="The arguments defining the " "model's architecture.",
|
|
205
|
+
)
|
|
206
|
+
group.add_argument(
|
|
207
|
+
"--hiddenEncoder",
|
|
208
|
+
type=int,
|
|
209
|
+
default=256,
|
|
210
|
+
help="Hidden dimension of the encoder network.",
|
|
211
|
+
)
|
|
212
|
+
group.add_argument(
|
|
213
|
+
"--hiddenGar",
|
|
214
|
+
type=int,
|
|
215
|
+
default=256,
|
|
216
|
+
help="Hidden dimension of the auto-regressive network",
|
|
217
|
+
)
|
|
218
|
+
group.add_argument(
|
|
219
|
+
"--nPredicts", type=int, default=12, help="Number of steps to predict."
|
|
220
|
+
)
|
|
221
|
+
group.add_argument(
|
|
222
|
+
"--negativeSamplingExt",
|
|
223
|
+
type=int,
|
|
224
|
+
default=128,
|
|
225
|
+
help="Number of negative samples to take.",
|
|
226
|
+
)
|
|
227
|
+
group.add_argument("--learningRate", type=float, default=2e-4)
|
|
228
|
+
group.add_argument(
|
|
229
|
+
"--schedulerStep",
|
|
230
|
+
type=int,
|
|
231
|
+
default=-1,
|
|
232
|
+
help="Step of the learning rate scheduler: at each "
|
|
233
|
+
"step the learning rate is divided by 2. Default: "
|
|
234
|
+
"no scheduler.",
|
|
235
|
+
)
|
|
236
|
+
group.add_argument(
|
|
237
|
+
"--schedulerRamp",
|
|
238
|
+
type=int,
|
|
239
|
+
default=None,
|
|
240
|
+
help="Enable a warm up phase for the learning rate: "
|
|
241
|
+
"adds a linear ramp of the given size.",
|
|
242
|
+
)
|
|
243
|
+
group.add_argument(
|
|
244
|
+
"--beta1",
|
|
245
|
+
type=float,
|
|
246
|
+
default=0.9,
|
|
247
|
+
help="Value of beta1 for the Adam optimizer",
|
|
248
|
+
)
|
|
249
|
+
group.add_argument(
|
|
250
|
+
"--beta2",
|
|
251
|
+
type=float,
|
|
252
|
+
default=0.999,
|
|
253
|
+
help="Value of beta2 for the Adam optimizer",
|
|
254
|
+
)
|
|
255
|
+
group.add_argument(
|
|
256
|
+
"--epsilon",
|
|
257
|
+
type=float,
|
|
258
|
+
default=1e-08,
|
|
259
|
+
help="Value of epsilon for the Adam optimizer",
|
|
260
|
+
)
|
|
261
|
+
group.add_argument(
|
|
262
|
+
"--sizeWindow",
|
|
263
|
+
type=int,
|
|
264
|
+
default=20480,
|
|
265
|
+
help="Number of frames to consider at each batch.",
|
|
266
|
+
)
|
|
267
|
+
group.add_argument(
|
|
268
|
+
"--nEpoch", type=int, default=200, help="Number of epoch to run"
|
|
269
|
+
)
|
|
270
|
+
group.add_argument(
|
|
271
|
+
"--samplingType",
|
|
272
|
+
type=str,
|
|
273
|
+
default="samespeaker",
|
|
274
|
+
choices=["samespeaker", "uniform", "samesequence", "sequential"],
|
|
275
|
+
help="How to sample the negative examples in the " "CPC loss.",
|
|
276
|
+
)
|
|
277
|
+
group.add_argument(
|
|
278
|
+
"--nLevelsPhone",
|
|
279
|
+
type=int,
|
|
280
|
+
default=1,
|
|
281
|
+
help="(Supervised mode only). Number of layers in "
|
|
282
|
+
"the phone classification network.",
|
|
283
|
+
)
|
|
284
|
+
group.add_argument(
|
|
285
|
+
"--cpc_mode",
|
|
286
|
+
type=str,
|
|
287
|
+
default=None,
|
|
288
|
+
choices=["reverse", "none"],
|
|
289
|
+
help="Some variations on CPC.",
|
|
290
|
+
)
|
|
291
|
+
group.add_argument(
|
|
292
|
+
"--encoder_type",
|
|
293
|
+
type=str,
|
|
294
|
+
choices=["cpc", "mfcc", "lfb"],
|
|
295
|
+
default="cpc",
|
|
296
|
+
help="Replace the encoder network by mfcc features "
|
|
297
|
+
"or learned filter banks",
|
|
298
|
+
)
|
|
299
|
+
group.add_argument(
|
|
300
|
+
"--normMode",
|
|
301
|
+
type=str,
|
|
302
|
+
default="layerNorm",
|
|
303
|
+
choices=["instanceNorm", "ID", "layerNorm", "batchNorm"],
|
|
304
|
+
help="Type of normalization to use in the encoder "
|
|
305
|
+
"network (default is layerNorm).",
|
|
306
|
+
)
|
|
307
|
+
group.add_argument(
|
|
308
|
+
"--onEncoder",
|
|
309
|
+
action="store_true",
|
|
310
|
+
help="(Supervised mode only) Perform the "
|
|
311
|
+
"classification on the encoder's output.",
|
|
312
|
+
)
|
|
313
|
+
group.add_argument(
|
|
314
|
+
"--random_seed", type=int, default=None, help="Set a specific random seed."
|
|
315
|
+
)
|
|
316
|
+
group.add_argument(
|
|
317
|
+
"--speakerEmbedding",
|
|
318
|
+
type=int,
|
|
319
|
+
default=0,
|
|
320
|
+
help="(Depreciated) Feed the prediction network with "
|
|
321
|
+
"speaker embeddings along with the usual sequence.",
|
|
322
|
+
)
|
|
323
|
+
group.add_argument(
|
|
324
|
+
"--arMode",
|
|
325
|
+
default="LSTM",
|
|
326
|
+
choices=["GRU", "LSTM", "RNN", "no_ar", "transformer"],
|
|
327
|
+
help="Architecture to use for the auto-regressive "
|
|
328
|
+
"network (default is lstm).",
|
|
329
|
+
)
|
|
330
|
+
group.add_argument(
|
|
331
|
+
"--nLevelsGRU",
|
|
332
|
+
type=int,
|
|
333
|
+
default=1,
|
|
334
|
+
help="Number of layers in the autoregressive network.",
|
|
335
|
+
)
|
|
336
|
+
group.add_argument(
|
|
337
|
+
"--rnnMode",
|
|
338
|
+
type=str,
|
|
339
|
+
default="transformer",
|
|
340
|
+
choices=[
|
|
341
|
+
"transformer",
|
|
342
|
+
"RNN",
|
|
343
|
+
"LSTM",
|
|
344
|
+
"linear",
|
|
345
|
+
"ffd",
|
|
346
|
+
"conv4",
|
|
347
|
+
"conv8",
|
|
348
|
+
"conv12",
|
|
349
|
+
],
|
|
350
|
+
help="Architecture to use for the prediction network",
|
|
351
|
+
)
|
|
352
|
+
group.add_argument(
|
|
353
|
+
"--dropout",
|
|
354
|
+
action="store_true",
|
|
355
|
+
help="Add a dropout layer at the output of the " "prediction network.",
|
|
356
|
+
)
|
|
357
|
+
group.add_argument(
|
|
358
|
+
"--abspos",
|
|
359
|
+
action="store_true",
|
|
360
|
+
help="If the prediction network is a transformer, "
|
|
361
|
+
"active to use absolute coordinates.",
|
|
362
|
+
)
|
|
363
|
+
return parser.parse_args([])
|
|
364
|
+
|
|
365
|
+
# from cpc.model import CPCModel as cpcmodel
|
|
366
|
+
# from cpc.cpc_default_config import get_default_cpc_config
|
|
367
|
+
# from cpc.feature_loader import getEncoder, getAR, loadArgs
|
|
368
|
+
# from cpc.feature_loader import loadArgs
|
|
369
|
+
|
|
370
|
+
locArgs = get_default_cpc_config()
|
|
371
|
+
|
|
372
|
+
if exists(checkpoint_cpc):
|
|
373
|
+
checkpoint = torch.load(checkpoint_cpc, map_location="cpu")
|
|
374
|
+
else:
|
|
375
|
+
checkpoint_url = "https://dl.fbaipublicfiles.com/librilight/CPC_checkpoints/60k_epoch4-d0f474de.pt"
|
|
376
|
+
checkpoint = torch.hub.load_state_dict_from_url(
|
|
377
|
+
checkpoint_url, progress=False, map_location="cpu"
|
|
378
|
+
)
|
|
379
|
+
makedirs(dirname(checkpoint_cpc))
|
|
380
|
+
torch.save(checkpoint, checkpoint_cpc)
|
|
381
|
+
|
|
382
|
+
temp = {"cpc": checkpoint_cpc}
|
|
383
|
+
loadArgs(locArgs, argparse.Namespace(**temp))
|
|
384
|
+
# encoderNet = getEncoder(locArgs)
|
|
385
|
+
encoderNet = CPCEncoder(locArgs.hiddenEncoder, locArgs.normMode)
|
|
386
|
+
# arNet = getAR(locArgs)
|
|
387
|
+
arNet = CPCAR(
|
|
388
|
+
locArgs.hiddenEncoder,
|
|
389
|
+
locArgs.hiddenGar,
|
|
390
|
+
locArgs.samplingType == "sequential",
|
|
391
|
+
locArgs.nLevelsGRU,
|
|
392
|
+
mode=locArgs.arMode,
|
|
393
|
+
reverse=locArgs.cpc_mode == "reverse",
|
|
394
|
+
)
|
|
395
|
+
# model = cpcmodel(encoderNet, arNet)
|
|
396
|
+
model = CPCModel(encoderNet, arNet)
|
|
397
|
+
|
|
398
|
+
# always load pretrained
|
|
399
|
+
if load_state_dict:
|
|
400
|
+
print("#" * 40)
|
|
401
|
+
print("Load pretrained CPC")
|
|
402
|
+
print("#" * 40)
|
|
403
|
+
model.load_state_dict(checkpoint["weights"], strict=False)
|
|
404
|
+
model.name = "cpc"
|
|
405
|
+
return model
|
|
406
|
+
|
|
407
|
+
|
|
408
|
+
class LayerNorm(nn.Module):
|
|
409
|
+
"""
|
|
410
|
+
Extending `nn.LayerNorm` by rearranging input dims to normalize over channel dimension in convnets.
|
|
411
|
+
|
|
412
|
+
The original `nn.LayerNorm` + 2 einops Rearrange is faster than custom Norm which calculated values directly on channel...
|
|
413
|
+
"""
|
|
414
|
+
|
|
415
|
+
def __init__(self, dim: int, rearrange_outputs: bool = True) -> None:
|
|
416
|
+
super().__init__()
|
|
417
|
+
self.ln = nn.LayerNorm(dim)
|
|
418
|
+
self.in_rearrange = Rearrange("b d t -> b t d")
|
|
419
|
+
if rearrange_outputs:
|
|
420
|
+
self.out_rearrange = Rearrange("b t d -> b d t")
|
|
421
|
+
else:
|
|
422
|
+
self.out_rearrange = nn.Identity()
|
|
423
|
+
|
|
424
|
+
def __repr__(self):
|
|
425
|
+
return str(self.ln)
|
|
426
|
+
|
|
427
|
+
def forward(self, x):
|
|
428
|
+
return self.out_rearrange(self.ln(self.in_rearrange(x)))
|
|
429
|
+
|
|
430
|
+
|
|
431
|
+
class CConv1d(nn.Conv1d):
|
|
432
|
+
"""source: https://github.com/pytorch/pytorch/issues/1333"""
|
|
433
|
+
|
|
434
|
+
def __init__(
|
|
435
|
+
self,
|
|
436
|
+
in_channels,
|
|
437
|
+
out_channels,
|
|
438
|
+
kernel_size,
|
|
439
|
+
stride=1,
|
|
440
|
+
dilation=1,
|
|
441
|
+
groups=1,
|
|
442
|
+
padding_value=0,
|
|
443
|
+
bias=True,
|
|
444
|
+
**kwargs,
|
|
445
|
+
):
|
|
446
|
+
super().__init__(
|
|
447
|
+
in_channels,
|
|
448
|
+
out_channels,
|
|
449
|
+
kernel_size=kernel_size,
|
|
450
|
+
stride=stride,
|
|
451
|
+
dilation=dilation,
|
|
452
|
+
groups=groups,
|
|
453
|
+
bias=bias,
|
|
454
|
+
**kwargs,
|
|
455
|
+
)
|
|
456
|
+
|
|
457
|
+
ks = kernel_size if isinstance(kernel_size, int) else kernel_size[0]
|
|
458
|
+
pad_dim1_pre = ks - 1
|
|
459
|
+
pad_dim1_post = 0
|
|
460
|
+
if dilation > 0:
|
|
461
|
+
pad_dim1_pre *= dilation
|
|
462
|
+
pad = (pad_dim1_pre, pad_dim1_post)
|
|
463
|
+
self.pad = nn.ConstantPad1d(padding=pad, value=padding_value)
|
|
464
|
+
|
|
465
|
+
def debug_weights(self, type="sum"):
|
|
466
|
+
w = 1.0
|
|
467
|
+
if type == "mean":
|
|
468
|
+
w = 1.0 / self.kernel_size[0]
|
|
469
|
+
|
|
470
|
+
elif type == "range":
|
|
471
|
+
k = self.kernel_size[0]
|
|
472
|
+
w = torch.arange(1, k + 1).float().pow(2)
|
|
473
|
+
w = w.repeat(self.out_channels, self.in_channels, 1)
|
|
474
|
+
print("w: ", w.shape)
|
|
475
|
+
self.weight.data = self.weight.data = w
|
|
476
|
+
if self.bias:
|
|
477
|
+
self.bias.data = self.bias.data.fill_(0.0)
|
|
478
|
+
return None
|
|
479
|
+
|
|
480
|
+
self.weight.data = self.weight.data.fill_(w)
|
|
481
|
+
if self.bias:
|
|
482
|
+
self.bias.data = self.bias.data.fill_(0.0)
|
|
483
|
+
|
|
484
|
+
def forward(self, input_):
|
|
485
|
+
# a =
|
|
486
|
+
# print(a[0, :, :6])
|
|
487
|
+
# print(a.shape)
|
|
488
|
+
# input("a")
|
|
489
|
+
# b =
|
|
490
|
+
# print(b[0, :, 0])
|
|
491
|
+
# print(b.shape)
|
|
492
|
+
# input("b")
|
|
493
|
+
return super().forward(self.pad(input_))
|
|
494
|
+
|
|
495
|
+
|
|
496
|
+
def get_cnn_layer(
|
|
497
|
+
dim: int,
|
|
498
|
+
kernel: List[int] = [5],
|
|
499
|
+
stride: List[int] = [2],
|
|
500
|
+
dilation: List[int] = [1],
|
|
501
|
+
activation: str = "GELU",
|
|
502
|
+
):
|
|
503
|
+
layers = []
|
|
504
|
+
layers.append(Rearrange("b t d -> b d t"))
|
|
505
|
+
for k, s, d in zip(kernel, stride, dilation):
|
|
506
|
+
#layers.append(CConv1d(dim, dim, kernel_size=k, stride=s, dilation=d))
|
|
507
|
+
layers.append(nn.Conv1d(dim, dim, kernel_size=k, stride=s, dilation=d))
|
|
508
|
+
layers.append(LayerNorm(dim))
|
|
509
|
+
layers.append(getattr(torch.nn, activation)())
|
|
510
|
+
layers.append(Rearrange("b d t -> b t d"))
|
|
511
|
+
return nn.Sequential(*layers)
|