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 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91
| import torch import torch.nn as nn
class EmoNet(nn.Module): """ EmoNet: 双流情感分类网络 Stream 1: 微多普勒图 → CNN提取空间特征 Stream 2: VTP序列 → BiLSTM提取时序特征 融合层: 拼接+MLP → 情绪分类 """ def __init__(self, n_classes: int = 6, n_doppler_bins: int = 33): super().__init__() self.doppler_cnn = nn.Sequential( nn.Conv2d(1, 32, 3, padding=1), nn.BatchNorm2d(32), nn.ReLU(), nn.MaxPool2d(2), nn.Conv2d(32, 64, 3, padding=1), nn.BatchNorm2d(64), nn.ReLU(), nn.MaxPool2d(2), nn.Conv2d(64, 128, 3, padding=1), nn.BatchNorm2d(128), nn.ReLU(), nn.AdaptiveAvgPool2d(1), ) self.vtp_lstm = nn.LSTM( input_size=1, hidden_size=64, num_layers=2, batch_first=True, bidirectional=True, dropout=0.3 ) self.classifier = nn.Sequential( nn.Linear(128 + 128, 64), nn.ReLU(), nn.Dropout(0.5), nn.Linear(64, n_classes) ) def forward(self, doppler_spec: torch.Tensor, vtp_seq: torch.Tensor) -> torch.Tensor: """ Args: doppler_spec: [B, 1, T, F] 微多普勒图 vtp_seq: [B, T, 1] 速度时间剖面 Returns: logits: [B, n_classes] """ dop_feat = self.doppler_cnn(doppler_spec) dop_feat = dop_feat.flatten(1) lstm_out, _ = self.vtp_lstm(vtp_seq) vtp_feat = lstm_out[:, -1, :] fused = torch.cat([dop_feat, vtp_feat], dim=1) logits = self.classifier(fused) return logits
if __name__ == "__main__": model = EmoNet(n_classes=6) doppler_input = torch.randn(4, 1, 90, 33) vtp_input = torch.randn(4, 90, 1) output = model(doppler_input, vtp_input) print(f"微多普勒输入: {doppler_input.shape}") print(f"VTP输入: {vtp_input.shape}") print(f"输出logits: {output.shape}") print(f"模型参数量: {sum(p.numel() for p in model.parameters()):,}")
|