File size: 2,279 Bytes
e1bc36a a2848e3 e1bc36a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 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 | import os
import sys
import torch
import torch.nn as nn
sys.path.append(os.getcwd())
from infer.lib.predictors.RMVPE.deepunet import DeepUnet, HPADeepUnet
N_MELS, N_CLASS = 128, 360
class BiGRU(nn.Module):
def __init__(
self,
input_features,
hidden_features,
num_layers
):
super(BiGRU, self).__init__()
self.gru = nn.GRU(
input_features,
hidden_features,
num_layers=num_layers,
batch_first=True,
bidirectional=True
)
def forward(self, x):
try:
return self.gru(x)[0]
except:
torch.backends.cudnn.enabled = False
return self.gru(x)[0]
class E2E(nn.Module):
def __init__(
self,
n_blocks,
n_gru,
kernel_size,
en_de_layers=5,
inter_layers=4,
in_channels=1,
en_out_channels=16,
hpa=False
):
super(E2E, self).__init__()
self.unet = (
HPADeepUnet(
in_channels=in_channels,
en_out_channels=en_out_channels,
base_channels=64,
hyperace_k=2,
hyperace_l=1,
num_hyperedges=16,
num_heads=4
)
) if hpa else (
DeepUnet(
kernel_size,
n_blocks,
en_de_layers,
inter_layers,
in_channels,
en_out_channels
)
)
self.cnn = nn.Conv2d(en_out_channels, 3, (3, 3), padding=(1, 1))
self.fc = (
nn.Sequential(
BiGRU(3 * 128, 256, n_gru),
nn.Linear(512, N_CLASS),
nn.Dropout(0.25),
nn.Sigmoid()
)
) if n_gru else (
nn.Sequential(
nn.Linear(3 * N_MELS, N_CLASS),
nn.Dropout(0.25),
nn.Sigmoid()
)
)
def forward(self, mel):
return self.fc(
self.cnn(
self.unet(
mel.transpose(-1, -2).unsqueeze(1)
)
).transpose(1, 2).flatten(-2)
) |