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 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168
| import torch import torch.nn as nn import torch.nn.functional as F
class ConvNeXt_MHSA_BiGRU(nn.Module): """ 分布式雷达连续动作识别框架 输入: 5个雷达节点的多普勒-时间频谱图 输出: 逐帧动作分类(9类) 架构: 1. ConvNeXt 编码器: 提取时频特征 2. MHSA 雷达视角融合: 自适应多节点融合 3. BiGRU 时序建模: 帧间动作转换 """ def __init__(self, n_radar=5, n_classes=9, hidden_dim=256): super().__init__() self.n_radar = n_radar self.encoder = ConvNeXtEncoder( in_channels=2, dims=[96, 192, 384, 768], depths=[3, 3, 9, 3] ) self.radar_mhsa = nn.MultiheadAttention( embed_dim=768, num_heads=8, batch_first=True, dropout=0.1 ) self.radar_norm = nn.LayerNorm(768) self.radar_dropout = nn.Dropout(0.2) self.time_mask = nn.Dropout(0.1) self.freq_mask = nn.Dropout(0.1) self.bigru = nn.GRU( input_size=768, hidden_size=hidden_dim, num_layers=2, batch_first=True, bidirectional=True, dropout=0.3 ) self.classifier = nn.Linear(hidden_dim * 2, n_classes) def forward(self, x): """ Args: x: (B, R, C, H, W) 雷达频谱图 B=batch, R=雷达数(5), C=2(实部+虚部), H=time, W=freq Returns: logits: (B, T, n_classes) 逐帧分类 """ B, R, C, H, W = x.shape x = x.view(B * R, C, H, W) if self.training: x = self.freq_mask(x) x = self.time_mask(x) features = self.encoder(x) T = features.shape[-1] features = features.view(B, R, -1, T) features = features.permute(0, 3, 1, 2) features_flat = features.reshape(B * T, R, -1) features_flat = self.radar_dropout(features_flat) attn_out, _ = self.radar_mhsa( features_flat, features_flat, features_flat ) fused = self.radar_norm(features_flat + attn_out) fused = fused.mean(dim=1) fused = fused.view(B, T, -1) gru_out, _ = self.bigru(fused) logits = self.classifier(gru_out) return logits
class ConvNeXtEncoder(nn.Module): """ConvNeXt 启发的频谱图编码器""" def __init__(self, in_channels=2, dims=[96, 192, 384, 768], depths=[3, 3, 9, 3]): super().__init__() self.stem = nn.Sequential( nn.Conv2d(in_channels, dims[0], 4, 4), nn.LayerNorm([dims[0]]), nn.GELU() ) self.stages = nn.ModuleList() for i in range(len(dims) - 1): stage = nn.Sequential( nn.Conv2d(dims[i], dims[i+1], 2, 2), *[ConvNeXtBlock(dims[i+1]) for _ in range(depths[i+1])] ) self.stages.append(stage) self.norm = nn.LayerNorm(dims[-1]) def forward(self, x): x = self.stem(x) for stage in self.stages: x = stage(x) x = x.mean(dim=(2, 3)) return x.unsqueeze(-1)
class ConvNeXtBlock(nn.Module): """ConvNeXt 残差块""" def __init__(self, dim): super().__init__() self.dwconv = nn.Conv2d(dim, dim, 7, 1, 3, groups=dim) self.norm = nn.LayerNorm(dim) self.pwconv1 = nn.Linear(dim, dim * 4) self.pwconv2 = nn.Linear(dim * 4, dim) self.act = nn.GELU() def forward(self, x): residual = x x = self.dwconv(x) x = x.permute(0, 2, 3, 1) x = self.norm(x) x = self.pwconv1(x) x = self.act(x) x = self.pwconv2(x) x = x.permute(0, 3, 1, 2) return residual + x
if __name__ == "__main__": model = ConvNeXt_MHSA_BiGRU(n_radar=5, n_classes=9) x = torch.randn(4, 5, 2, 64, 128) logits = model(x) print(f"Input: {x.shape}") print(f"Output: {logits.shape}") print(f"Parameters: {sum(p.numel() for p in model.parameters()):,}")
|