yanggeng1995 / multi-band-wavernn Goto Github PK

   for i, (x, y, m) in enumerate(train_set, 1):
        x, m, y = x.to(device), m.to(device), y.to(device)  # x/y: (Batch, sub_bands, T)

            y_hat = model(x, m)  # (Batch, T, num_classes, sub_bands)

            if model.mode == 'RAW':
                y_hat0 = y_hat[:, :, :, 0].transpose(1,2).unsqueeze(-1)  # (Batch, num_classes, T, 1)
                y_hat1 = y_hat[:, :, :, 1].transpose(1,2).unsqueeze(-1)
                y_hat2 = y_hat[:, :, :, 2].transpose(1,2).unsqueeze(-1)
                y_hat3 = y_hat[:, :, :, 3].transpose(1,2).unsqueeze(-1)

            elif model.mode == 'MOL':
                y0 = y0.float()
                y1 = y1.float()
                y2 = y2.float()
                y3 = y3.float()

            loss = loss_func(y_hat0, y0) + loss_func(y_hat1, y1) + loss_func(y_hat2, y2) + loss_func(y_hat3, y3)

def forward(self, x, mels):  # x: (Batch, Subband, T)
    device = next(self.parameters()).device  # use same device as parameters

    # Although we `_flatten_parameters()` on init, when using DataParallel
    # the model gets replicated, making it no longer guaranteed that the
    # weights are contiguous in GPU memory. Hence, we must call it again
    self._flatten_parameters()

    self.step += 1
    bsize = x.size(0)
    h1 = torch.zeros(1, bsize, self.rnn_dims, device=device)
    h2 = torch.zeros(1, bsize, self.rnn_dims, device=device)
    mels, aux = self.upsample(mels)

    aux_idx = [self.aux_dims * i for i in range(5)]
    a1 = aux[:, :, aux_idx[0]:aux_idx[1]]
    a2 = aux[:, :, aux_idx[1]:aux_idx[2]]
    a3 = aux[:, :, aux_idx[2]:aux_idx[3]]
    a4 = aux[:, :, aux_idx[3]:aux_idx[4]]
    # x = torch.cat([x.unsqueeze(-1), mels, a1], dim=2)

    x = torch.cat([x.transpose(1,2), mels, a1], dim=2)  # (batch,T,4)  (batch,T,80) (batch,T,32)
    x = self.I(x)  # (batch, T, 116) -> # (batch, T, 512)
    res = x
    # x, _ = self.rnn1(x, h1)
    x, _ = self.rnn1(x)   # 不加入隐藏层-Begee  # (batch, T, 512) -> (batch, T, 512)
    x = x + res    # (batch, T, 512)
    res = x
    x = torch.cat([x, a2], dim=2)  # (batch, T, 512) -> (batch, T, 512+128)
    # x, _ = self.rnn2(x, h2)
    x, _ = self.rnn2(x)  # 不加入隐藏层-Begee  (batch, T, 512+128) -> (batch, T, 512)

    x = x + res
    x = torch.cat([x, a3], dim=2) # (batch, T, 512+128)
    x = F.relu(self.fc1(x)) # (batch, T, 512+128) -> (batch, T, 512)

    x = torch.cat([x, a4], dim=2)  # (batch, T, 512+128)
    x = F.relu(self.fc2(x))  # (batch, T, 512+128) -> (batch, T, 512)

    out0 = self.fc30(x).unsqueeze(-1) # (batch, T, 512) -> (batch, T, 512)
    out1 = self.fc31(x).unsqueeze(-1)
    out2 = self.fc32(x).unsqueeze(-1)
    out3 = self.fc33(x).unsqueeze(-1)
    out = torch.cat([out0,out1,out2,out3], dim=3)  # (B, T, num_classes, sub_band)
    return out