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 ADDED
@@ -0,0 +1,12 @@
1
+
2
+ vap_main/output_offline.txt
3
+ /rvap.egg-info
4
+ all.bat
5
+ *.bat
6
+ /train/lightning_logs
7
+ *.onnx
8
+ .vscode/
9
+ tools/output_offline.txt
10
+ __pycache__/
11
+ tflite_vap/
12
+ tfjs_vap/
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,4 @@
1
+ from maai.model import Maai
2
+ import maai.input as MaaiInput
3
+ import maai.output as MaaiOutput
4
+ from maai.util import get_available_models
@@ -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)