ASLP-lab/DiffRhythm

1from einops import rearrange
2import numpy as np
3import torch
4import torch.nn as nn
5
6
7class Generator(nn.Module):
8
9    def __init__(self,
10                 in_features,
11                 ffd_hidden_size,
12                 num_classes,
13                 attn_layer_num,
14                 
15                 ):
16        super(Generator, self).__init__()
17        
18        self.attn = nn.ModuleList(
19            [
20                nn.MultiheadAttention(
21                    embed_dim=in_features,
22                    num_heads=8,
23                    dropout=0.2,
24                    batch_first=True,
25                )
26                for _ in range(attn_layer_num)
27            ]
28        )
29        
30        self.ffd = nn.Sequential(
31            nn.Linear(in_features, ffd_hidden_size),
32            nn.ReLU(),
33            nn.Linear(ffd_hidden_size, in_features)
34        )
35        
36        self.dropout = nn.Dropout(0.2)
37        
38        self.fc =  nn.Linear(in_features * 2, num_classes)
39        
40        self.proj = nn.Tanh()
41        
42
43    def forward(self, ssl_feature, judge_id=None):
44        '''
45        ssl_feature: [B, T, D]   
46        output: [B, num_classes]
47        '''
48        
49        B, T, D = ssl_feature.shape
50        
51        ssl_feature = self.ffd(ssl_feature)
52        
53        tmp_ssl_feature = ssl_feature
54        
55        for attn in self.attn:
56            tmp_ssl_feature, _ = attn(tmp_ssl_feature, tmp_ssl_feature, tmp_ssl_feature)
57    
58        ssl_feature = self.dropout(torch.concat([torch.mean(tmp_ssl_feature, dim=1), torch.max(ssl_feature, dim=1)[0]], dim=1))  # B, 2D
59        
60        x = self.fc(ssl_feature)  # B, num_classes
61        
62        x = self.proj(x) * 2.0 + 3
63        
64        return x
65    
66    
67