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
| class DriverFacialEmotionNet(nn.Module): """ 驾驶员面部情感识别网络 ======================= 论文 Figure 2: 整体架构 架构: 1. Stem: 7x7 Conv → BN → ReLU → MaxPool 2. Stage 1: 3× EnhancedResidualBlock (64) 3. Stage 2: 4× EnhancedResidualBlock (128) 4. Stage 3: 6× EnhancedResidualBlock (256) 5. Stage 4: 3× EnhancedResidualBlock (512) 6. Head: GlobalAvgPool → FC → Softmax 输入: 224×224 RGB 面部图像 输出: 7 类情感概率 Args: num_classes: 情感类别数 wcsa_layers: 哪些 stage 使用 WCSA """ def __init__( self, num_classes: int = 7, wcsa_layers: tuple = (True, True, True, True) ): super().__init__() self.stem = nn.Sequential( nn.Conv2d(3, 64, 7, 2, 3, bias=False), nn.BatchNorm2d(64), nn.ReLU(), nn.MaxPool2d(3, 2, 1) ) self.stage1 = nn.Sequential( EnhancedResidualBlock(64, 64, use_wcsa=wcsa_layers[0]), EnhancedResidualBlock(64, 64, use_wcsa=wcsa_layers[0]), EnhancedResidualBlock(64, 64, use_wcsa=wcsa_layers[0]) ) self.stage2 = nn.Sequential( EnhancedResidualBlock(64, 128, stride=2, use_wcsa=wcsa_layers[1]), EnhancedResidualBlock(128, 128, use_wcsa=wcsa_layers[1]), EnhancedResidualBlock(128, 128, use_wcsa=wcsa_layers[1]), EnhancedResidualBlock(128, 128, use_wcsa=wcsa_layers[1]) ) self.stage3 = nn.Sequential( EnhancedResidualBlock(128, 256, stride=2, use_wcsa=wcsa_layers[2]), EnhancedResidualBlock(256, 256, use_wcsa=wcsa_layers[2]), EnhancedResidualBlock(256, 256, use_wcsa=wcsa_layers[2]), EnhancedResidualBlock(256, 256, use_wcsa=wcsa_layers[2]), EnhancedResidualBlock(256, 256, use_wcsa=wcsa_layers[2]), EnhancedResidualBlock(256, 256, use_wcsa=wcsa_layers[2]) ) self.stage4 = nn.Sequential( EnhancedResidualBlock(256, 512, stride=2, use_wcsa=wcsa_layers[3]), EnhancedResidualBlock(512, 512, use_wcsa=wcsa_layers[3]), EnhancedResidualBlock(512, 512, use_wcsa=wcsa_layers[3]) ) self.head = nn.Sequential( nn.AdaptiveAvgPool2d(1), nn.Flatten(), nn.Dropout(0.3), nn.Linear(512, num_classes) ) def forward(self, x: torch.Tensor) -> torch.Tensor: x = self.stem(x) x = self.stage1(x) x = self.stage2(x) x = self.stage3(x) x = self.stage4(x) x = self.head(x) return x
if __name__ == "__main__": model = DriverFacialEmotionNet(num_classes=7) x = torch.randn(4, 3, 224, 224) out = model(x) total_params = sum(p.numel() for p in model.parameters()) print(f"驾驶员面部情感识别网络") print(f"输入: {x.shape}") print(f"输出: {out.shape}") print(f"总参数: {total_params:,} ({total_params/1e6:.2f}M)") model.eval() with torch.no_grad(): import time start = time.time() for _ in range(100): _ = model(x) elapsed = (time.time() - start) / 100 * 1000 print(f"推理延迟: {elapsed:.1f}ms (batch=4)") print(f"单帧延迟: {elapsed/4:.1f}ms")
|