-
Notifications
You must be signed in to change notification settings - Fork 26
/
_twin_net.py
53 lines (38 loc) · 1.29 KB
/
_twin_net.py
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
#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""The TwinNet of the Masker.
"""
from collections import namedtuple
from torch.nn import Module
from modules import _twin_rnn_dec, _fnn
__author__ = 'Konstantinos Drossos -- Tampere University'
__docformat__ = 'reStructuredText'
__all__ = ['TwinNet']
class TwinNet(Module):
def __init__(self, rnn_dec_input_dim, original_input_dim, context_length):
super(TwinNet, self).__init__()
self.rnn_dec = _twin_rnn_dec.TwinRNNDec(
input_dim=rnn_dec_input_dim
)
self.fnn = _fnn.FNNMasker(
input_dim=rnn_dec_input_dim,
output_dim=original_input_dim,
context_length=context_length
)
self.output = namedtuple(
'twin_net_output',
['h_dec_twin', 'v_j_filt_prime_twin']
)
def forward(self, h_enc, x):
"""The forward pass of the TwinNet.
:param h_enc: The input to the TwinNet.
:type h_enc: torch.Tensor
:param x: The original input to the non-twin\
counterpart.
:type x: torch.Tensor
:return: The output of the TwinNet.
:rtype: torch.Tensor
"""
h_dec_twin = self.rnn_dec(h_enc)
return self.output(h_dec_twin, self.fnn(h_dec_twin, x))
# EOF