Feature_Extraction in Conv-Tasnet

username: Kuldeep_Rana
cooked: <p>Hello, I want to extract the feature from our conv_tanset model and use these features for speaker diarization. Here I am uploading my model code. Thank you very much.</p>
<p>class ConvTasNet(nn.Module):<br>
def <strong>init</strong>(self,<br>
L=16,<br>
N=512,<br>
X=8,<br>
R=3,<br>
B=128,<br>
Sc=128,<br>
H=512,<br>
P=3,<br>
norm=“cLN”,<br>
num_spks=2,<br>
non_linear=“sigmoid”,<br>
causal=False):<br>
super(ConvTasNet, self).<strong>init</strong>()<br>
supported_nonlinear = {<br>
“relu”: F.relu,<br>
“sigmoid”: th.sigmoid,<br>
“softmax”: F.softmax<br>
}<br>
if non_linear not in supported_nonlinear:<br>
raise RuntimeError(“Unsupported non-linear function: {}”,<br>
format(non_linear))<br>
self.non_linear_type = non_linear<br>
self.non_linear = supported_nonlinear[non_linear]<br>
# n x S =&gt; n x N x T, S = 4s*8000 = 32000<br>
self.encoder_1d = Conv1D(1, N, L, stride=L // 2, padding=0)<br>
# keep T not change<br>
# T = int((xlen - L) / (L // 2)) + 1<br>
# before repeat blocks, always cLN<br>
self.ln = ChannelWiseLayerNorm(N)<br>
# n x N x T =&gt; n x B x T<br>
self.proj = Conv1D(N, B, 1)<br>
# repeat blocks<br>
# n x B x T =&gt; n x B x T<br>
self.repeats = self._build_repeats(<br>
R,<br>
X,<br>
Sc=Sc,<br>
in_channels=B,<br>
conv_channels=H,<br>
kernel_size=P,<br>
norm=norm,<br>
causal=causal)<br>
self.PRelu = nn.PReLU()<br>
# output 1x1 conv<br>
# n x B x T =&gt; n x N x T<br>
# NOTE: using ModuleList not python list<br>
# self.conv1x1_2 = th.nn.ModuleList(<br>
#     [Conv1D(B, N, 1) for _ in range(num_spks)])<br>
# n x Sc x T =&gt; n x 2N x T<br>
self.mask = Conv1D(Sc, num_spks * N, 1)<br>
# using ConvTrans1D: n x N x T =&gt; n x 1 x To<br>
# To = (T - 1) * L // 2 + L<br>
self.decoder_1d = ConvTrans1D(<br>
N, 1, kernel_size=L, stride=L // 2, bias=True)<br>
self.num_spks = num_spks<br>
self.R = R <span class="hashtag">#numbers</span> of repeat<br>
self.X = X <span class="hashtag">#numbers</span> of Conv1Dblock in each repeat</p>
<pre><code>def _build_blocks(self, num_blocks, **block_kwargs):
    """
    Build Conv1D block
    """
    blocks = [
        Conv1DBlock(**block_kwargs, dilation=(2**b))
        for b in range(num_blocks)
    ]
    return nn.Sequential(*blocks)

def _build_repeats(self, num_repeats, num_blocks, **block_kwargs):
    """
    Build Conv1D block repeats
    """
    repeats = [
        self._build_blocks(num_blocks, **block_kwargs)
        for r in range(num_repeats)
    ]
    return nn.Sequential(*repeats)

def forward(self, x):
    if x.dim() &gt;= 3:
        raise RuntimeError(
            "{} accept 1/2D tensor as input, but got {:d}".format(
                self.__name__, x.dim()))
    # when inference, only one utt
    if x.dim() == 1:
        x = th.unsqueeze(x, 0)
    #encoder
    # n x 1 x S =&gt; n x N x T
    w = F.relu(self.encoder_1d(x))
    
    #Seperation
    #   LayerNorm &amp; 1X1 Conv
    # n x B x T
    y = self.proj(self.ln(w))
    
    #TCN
    # n x B x T
    skip_connection = 0
    for i in range(self.R):
        for j in range(self.X):
            skip, y = self.repeats[i][j](y)
            skip_connection = skip_connection + skip
    
    y = self.PRelu(skip_connection)
    print(y.size())
    # n x 2N x T
    e = th.chunk(self.mask(y), self.num_spks, 1)
    # print(e.size())
    #e = create_feature_extractor(e)
    print("feature extraction", e)
    # n x N x T
    if self.non_linear_type == "softmax":
        m = self.non_linear(th.stack(e, dim=0), dim=0)
    else:
        m = self.non_linear(th.stack(e, dim=0))
    # spks x [n x N x T]
    s = [w * m[n] for n in range(self.num_spks)]
    # spks x n x S
    return [self.decoder_1d(x, squeeze=True) for x in s]</code></pre>