Spaces:
Sleeping
Sleeping
| import torch | |
| import torch.nn as nn | |
| from .sinc_conv import TimeSincExtractor, FreqSincExtractor | |
| from .patchify import Patchify | |
| from .csp_tiny_layer import CSPTinyLayer | |
| class TinyVAD(nn.Module): | |
| def __init__(self, in_channels, hidden_channels, out_channels, patch_size, num_blocks, sinc_conv, ssm): | |
| super(TinyVAD, self).__init__() | |
| self.sinc_conv = sinc_conv | |
| if self.sinc_conv: | |
| # self.extractor = TimeSincExtractor(out_channels=64, kernel_size=101, range_constraint=True, stride=2) | |
| self.extractor = FreqSincExtractor(out_channels=64, kernel_size=101, range_constraint=True, stride=2) | |
| self.patchify = Patchify(in_channels, hidden_channels, patch_size) | |
| self.csp_tiny_layer1 = CSPTinyLayer(hidden_channels, hidden_channels, num_blocks, ssm) | |
| self.csp_tiny_layer2 = CSPTinyLayer(hidden_channels, hidden_channels, num_blocks, ssm) | |
| self.csp_tiny_layer3 = CSPTinyLayer(hidden_channels, out_channels, num_blocks, ssm) | |
| self.avg_pool = nn.AdaptiveAvgPool2d(1) | |
| self.classifier = nn.Sequential( | |
| nn.Linear(out_channels, 1), | |
| # nn.Sigmoid() | |
| ) | |
| def forward(self, x): | |
| if self.sinc_conv: | |
| x = self.extractor(x, None) | |
| x = x[0] # Untuple | |
| x = self.patchify(x) | |
| x = self.csp_tiny_layer1(x) | |
| x = self.csp_tiny_layer2(x) | |
| x = self.csp_tiny_layer3(x) | |
| x = self.avg_pool(x).view(x.size(0), -1) | |
| x = self.classifier(x) | |
| return x | |
| def predict(self, inputs): | |
| logits = self.forward(inputs) | |
| probs = torch.sigmoid(logits) | |
| return probs | |
| if __name__ == "__main__": | |
| device = torch.device("cuda" if torch.cuda.is_available() else "cpu") | |
| print(f"Using device: {device}") | |
| model = TinyVAD(1, 32, 64, 2, 2, False, False).to(device) | |
| print(model) | |
| dummy_input = torch.randn(1, 1, 64, 16).to(device) | |
| output = model(dummy_input) | |
| print(output) | |