feat(actpred): 动作预测模块入库（PyTorch 训练/预测/ONNX 导出）

actpred/model.py

@@ -0,0 +1,117 @@
+import torch
+import torch.nn as nn
+from config import DataConfig, ModelConfig
+
+
+class SimpleCNN(nn.Module):
+    def __init__(self, feature_dim: int = 128):
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.Conv2d(3, 16, 3, stride=2, padding=1),
+            nn.BatchNorm2d(16),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(16, 32, 3, stride=2, padding=1),
+            nn.BatchNorm2d(32),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(32, 64, 3, stride=2, padding=1),
+            nn.BatchNorm2d(64),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(64, feature_dim, 3, stride=2, padding=1),
+            nn.BatchNorm2d(feature_dim),
+            nn.ReLU(inplace=True),
+            nn.AdaptiveAvgPool2d(1),
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.net(x).flatten(1)
+
+
+class MobileNetEncoder(nn.Module):
+    def __init__(self, feature_dim: int = 128):
+        super().__init__()
+        from torchvision.models import mobilenet_v2, MobileNet_V2_Weights
+        backbone = mobilenet_v2(weights=MobileNet_V2_Weights.DEFAULT)
+        self.features = backbone.features
+        self.pool = nn.AdaptiveAvgPool2d(1)
+        self.proj = nn.Linear(1280, feature_dim)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.features(x)
+        x = self.pool(x).flatten(1)
+        return self.proj(x)
+
+
+class MLPHistoryEncoder(nn.Module):
+    def __init__(self, input_dim: int, hidden_dim: int = 64, dropout: float = 0.1):
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.Linear(input_dim, hidden_dim),
+            nn.ReLU(inplace=True),
+            nn.Dropout(dropout),
+            nn.Linear(hidden_dim, hidden_dim),
+            nn.ReLU(inplace=True),
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.net(x)
+
+
+class GRUHistoryEncoder(nn.Module):
+    def __init__(self, input_dim: int, hidden_dim: int = 64, dropout: float = 0.1):
+        super().__init__()
+        self.gru = nn.GRU(input_size=input_dim, hidden_size=hidden_dim,
+                          num_layers=1, batch_first=True)
+        self.dropout = nn.Dropout(dropout)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        _, h = self.gru(x)
+        return self.dropout(h.squeeze(0))
+
+
+class FusionHead(nn.Module):
+    def __init__(self, input_dim: int, hidden_dim: int, output_dim: int, dropout: float = 0.1):
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.Linear(input_dim, hidden_dim),
+            nn.ReLU(inplace=True),
+            nn.Dropout(dropout),
+            nn.Linear(hidden_dim, hidden_dim // 2),
+            nn.ReLU(inplace=True),
+            nn.Linear(hidden_dim // 2, output_dim),
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.net(x)
+
+
+class ActionPredictor(nn.Module):
+    def __init__(self, model_cfg: ModelConfig, data_cfg: DataConfig):
+        super().__init__()
+        if model_cfg.cnn_backbone == "mobilenet":
+            self.image_encoder = MobileNetEncoder(model_cfg.cnn_feature_dim)
+        else:
+            self.image_encoder = SimpleCNN(model_cfg.cnn_feature_dim)
+
+        history_input_dim = data_cfg.history_len * data_cfg.action_dim
+        if model_cfg.history_encoder == "gru":
+            self.action_encoder = GRUHistoryEncoder(
+                data_cfg.action_dim, model_cfg.action_hidden_dim, model_cfg.dropout)
+        else:
+            self.action_encoder = MLPHistoryEncoder(
+                history_input_dim, model_cfg.action_hidden_dim, model_cfg.dropout)
+
+        fused_input = model_cfg.cnn_feature_dim + model_cfg.action_hidden_dim
+        self.head = FusionHead(
+            fused_input, model_cfg.fused_hidden_dim, data_cfg.action_dim, model_cfg.dropout)
+
+        self._use_gru = model_cfg.history_encoder == "gru"
+        self._history_len = data_cfg.history_len
+        self._action_dim = data_cfg.action_dim
+
+    def forward(self, image: torch.Tensor, action_history: torch.Tensor) -> torch.Tensor:
+        img_feat = self.image_encoder(image)
+        if self._use_gru:
+            act_feat = self.action_encoder(action_history)
+        else:
+            act_feat = self.action_encoder(action_history.flatten(1))
+        return self.head(torch.cat([img_feat, act_feat], dim=-1))