| from __future__ import annotations |
|
|
| import torch |
| from torch import nn |
| from torch.nn import functional as F |
|
|
|
|
| class PocketMoE(nn.Module): |
| def __init__(self, experts: int = 4, top_k: int = 2) -> None: |
| super().__init__() |
| self.expert_count = experts |
| self.top_k = top_k |
| self.encoder = nn.Sequential( |
| nn.Linear(64, 32), |
| nn.GELU(), |
| ) |
| self.router = nn.Linear(32, experts) |
| self.experts = nn.ModuleList( |
| [ |
| nn.Sequential( |
| nn.Linear(32, 16), |
| nn.GELU(), |
| nn.Linear(16, 10), |
| ) |
| for _ in range(experts) |
| ] |
| ) |
|
|
| def forward( |
| self, |
| pixels: torch.Tensor, |
| ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: |
| hidden = self.encoder(pixels) |
| router_probabilities = F.softmax(self.router(hidden), dim=1) |
| top_probabilities, top_indices = router_probabilities.topk( |
| self.top_k, |
| dim=1, |
| ) |
| sparse_weights = torch.zeros_like(router_probabilities).scatter( |
| 1, |
| top_indices, |
| top_probabilities, |
| ) |
| sparse_weights = sparse_weights / sparse_weights.sum(dim=1, keepdim=True) |
| expert_logits = torch.stack( |
| [expert(hidden) for expert in self.experts], |
| dim=1, |
| ) |
| logits = (expert_logits * sparse_weights.unsqueeze(-1)).sum(dim=1) |
| return logits, router_probabilities, sparse_weights |
|
|
|
|
| class DenseControl(nn.Module): |
| def __init__(self) -> None: |
| super().__init__() |
| self.network = nn.Sequential( |
| nn.Linear(64, 48), |
| nn.GELU(), |
| nn.Linear(48, 40), |
| nn.GELU(), |
| nn.Linear(40, 10), |
| ) |
|
|
| def forward(self, pixels: torch.Tensor) -> torch.Tensor: |
| return self.network(pixels) |
|
|
|
|
| def parameter_count(model: nn.Module) -> int: |
| return sum(parameter.numel() for parameter in model.parameters()) |
|
|
|
|
| def active_parameter_count(model: PocketMoE) -> int: |
| shared = sum(parameter.numel() for parameter in model.encoder.parameters()) |
| router = sum(parameter.numel() for parameter in model.router.parameters()) |
| experts = sorted( |
| [ |
| sum(parameter.numel() for parameter in expert.parameters()) |
| for expert in model.experts |
| ], |
| reverse=True, |
| ) |
| return shared + router + sum(experts[: model.top_k]) |
|
|
