118 lines
4.2 KiB
Python
118 lines
4.2 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
|
|
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))
|