qing
3d03f7a1ce
审计报告: docs/reviews/actpred_CODE_AUDIT_20260808.md - P0(2项): episode索引错位(dataset.py)、推理归一化静默失败(predict.py) - P1(6项): 空训练集校验、图像自然排序、空样本防御、形状校验、 MobileNet下载错误提示、stats路径校验 - P2(8项): README示例修正、公共属性、glob优化、版本约束、 单元测试(17项)、GPU安装说明、BooleanOptionalAction、ONNX验证 新增文件: actpred/test_actpred.py (17项单元测试, 全部通过) 更新文档: README.md配置表扩充、INDEX.md、审计报告修复总结
127 lines
4.7 KiB
Python
127 lines
4.7 KiB
Python
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
|
|
try:
|
|
backbone = mobilenet_v2(weights=MobileNet_V2_Weights.DEFAULT)
|
|
except Exception as e:
|
|
raise RuntimeError(
|
|
"MobileNetV2 预训练权重下载失败，可能是网络受限。"
|
|
"请在联网环境中预下载权重"
|
|
"（运行 torchvision.models.mobilenet_v2(weights='IMAGENET1K_V1')），"
|
|
"并将 ~/.cache/torch/hub/checkpoints/ 下的权重文件"
|
|
"拷贝到目标机器同路径。"
|
|
) from e
|
|
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))
|