| | import torch
|
| | import torch.nn as nn
|
| |
|
| |
|
| | class PitchRNN(nn.Module):
|
| | def __init__(self, n_mels, hidden_size):
|
| | super(PitchRNN, self).__init__()
|
| |
|
| | self.sp_linear = nn.Sequential(nn.Conv1d(n_mels, hidden_size*2, kernel_size=1),
|
| | nn.SiLU(),
|
| | nn.Conv1d(hidden_size*2, hidden_size, kernel_size=1),
|
| | nn.SiLU(),)
|
| |
|
| | self.midi_linear = nn.Sequential(nn.Conv1d(1, hidden_size*2, kernel_size=1),
|
| | nn.SiLU(),
|
| | nn.Conv1d(hidden_size*2, hidden_size, kernel_size=1),
|
| | nn.SiLU(),)
|
| |
|
| | self.hidden_size = hidden_size
|
| |
|
| | self.rnn = nn.GRU(input_size=hidden_size*2,
|
| | hidden_size=hidden_size,
|
| | num_layers=2,
|
| | batch_first=True,
|
| | bidirectional=True)
|
| |
|
| |
|
| | self.linear = nn.Sequential(nn.Linear(2*hidden_size, hidden_size),
|
| | nn.SiLU(),
|
| | nn.Linear(hidden_size, 1))
|
| |
|
| | def forward(self, midi, sp):
|
| | midi = midi.unsqueeze(1)
|
| | midi = self.midi_linear(midi)
|
| | sp = self.sp_linear(sp)
|
| |
|
| | x = torch.cat([midi, sp], dim=1)
|
| | x = torch.transpose(x, 1, 2)
|
| | x, _ = self.rnn(x)
|
| |
|
| |
|
| | x = self.linear(x)
|
| |
|
| | return x.squeeze(-1)
|
| |
|
| |
|
| | if __name__ == '__main__':
|
| |
|
| | model = PitchRNN(100, 256)
|
| |
|
| | x = torch.rand((4, 128))
|
| | t = torch.randint(0, 1000, (1, )).long()
|
| | sp = torch.rand((4, 100, 128))
|
| | midi = torch.rand((4, 128))
|
| |
|
| | y = model(midi, sp) |
