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
| import torch import torch.nn as nn
class DistributedRadarHAR(nn.Module): """分布式雷达人体活动识别""" def __init__(self, num_sensors=5, num_classes=10, hidden_dim=256, num_heads=8): super().__init__() self.sensor_encoders = nn.ModuleList([ ConvNeXtEncoder(out_dim=hidden_dim) for _ in range(num_sensors) ]) self.sensor_fusion = nn.MultiheadAttention( hidden_dim, num_heads, batch_first=True ) self.mhsa = nn.MultiheadAttention( hidden_dim, num_heads, batch_first=True ) self.bigru = nn.GRU( hidden_dim, hidden_dim // 2, num_layers=2, batch_first=True, bidirectional=True ) self.classifier = nn.Sequential( nn.Linear(hidden_dim, 128), nn.GELU(), nn.Dropout(0.2), nn.Linear(128, num_classes) ) def forward(self, radar_sequence): """ Args: radar_sequence: (B, T, S, 1, H, W) T=time, S=sensors, H×W=range-doppler map """ B, T, S, _, H, W = radar_sequence.shape sensor_feats = [] for s in range(S): feat = self.sensor_encoders[s]( radar_sequence[:, :, s] ) sensor_feats.append(feat) fused = torch.stack(sensor_feats, dim=2) fused_flat = fused.reshape(B * T, S, -1) attn_out, _ = self.sensor_fusion( fused_flat, fused_flat, fused_flat ) fused = attn_out.mean(dim=1).reshape(B, T, -1) normed = nn.LayerNorm(fused.size(-1)).to(fused.device)(fused) mhsa_out, _ = self.mhsa(normed, normed, normed) fused = fused + mhsa_out gru_out, _ = self.bigru(fused) out = gru_out[:, -1, :] return self.classifier(out)
class ConvNeXtEncoder(nn.Module): """简化 ConvNeXt 编码器""" def __init__(self, out_dim=256): super().__init__() self.encoder = nn.Sequential( nn.Conv2d(1, 32, 7, stride=2, padding=3), nn.BatchNorm2d(32), nn.GELU(), nn.Conv2d(32, 64, 3, stride=2, padding=1), nn.BatchNorm2d(64), nn.GELU(), nn.Conv2d(64, 128, 3, stride=2, padding=1), nn.BatchNorm2d(128), nn.GELU(), nn.AdaptiveAvgPool2d(1), nn.Flatten(), nn.Linear(128, out_dim) ) def forward(self, x): B, T = x.shape[:2] x = x.reshape(B * T, *x.shape[2:]) feat = self.encoder(x) return feat.reshape(B, T, -1)
if __name__ == "__main__": model = DistributedRadarHAR(num_sensors=5, num_classes=10) x = torch.randn(2, 30, 5, 1, 64, 64) out = model(x) print(f"Input: {x.shape}") print(f"Output: {out.shape}") print(f"Params: {sum(p.numel() for p in model.parameters()):,}")
|