Add nano and tiny model presets

This commit is contained in:
2026-07-23 15:01:33 +08:00
parent 06f29c0f0a
commit 6f7b5be405
3 changed files with 38 additions and 29 deletions

View File

@@ -15,7 +15,7 @@ from backbones import (
from targets import PAD_IDX
EVENT_TRAJECTORY_ARCHITECTURE = "event_trajectory_shared_v1"
EVENT_TRAJECTORY_ARCHITECTURE = "event_trajectory_shared_v2"
@dataclass(frozen=True)
@@ -33,7 +33,8 @@ class EventTrajectoryModelSize:
MODEL_SIZE_PRESETS = {
"nano": EventTrajectoryModelSize(d_model=256, n_trajectory=8),
"nano": EventTrajectoryModelSize(d_model=120, n_trajectory=6),
"tiny": EventTrajectoryModelSize(d_model=256, n_trajectory=8),
"small": EventTrajectoryModelSize(d_model=512, n_trajectory=8),
"medium": EventTrajectoryModelSize(d_model=768, n_trajectory=12),
"huge": EventTrajectoryModelSize(d_model=1024, n_trajectory=16),