Converts a model variant string into a dictionary of hyperparameters.
(variant: Optional[str])
| 478 | |
| 479 | |
| 480 | def decode_variant(variant: Optional[str]) -> dict: |
| 481 | """Converts a model variant string into a dictionary of hyperparameters.""" |
| 482 | if not variant: |
| 483 | return {} |
| 484 | |
| 485 | # Hyperparameters from text transformer models, often similar to ViT sizes. |
| 486 | return { |
| 487 | "width": {"Ti": 192, "S": 384, "M": 512, "B": 512, "L": 768, "H": 1024, "G": 1664}[variant], |
| 488 | "depth": {"Ti": 12, "S": 12, "M": 12, "B": 12, "L": 12, "H": 24, "G": 48}[variant], |
| 489 | "mlp_dim": {"Ti": 768, "S": 1536, "M": 2048, "B": 2048, "L": 3072, "H": 4096, "G": 8192}[variant], |
| 490 | "num_heads": {"Ti": 3, "S": 6, "M": 8, "B": 8, "L": 12, "H": 16, "G": 16}[variant], |
| 491 | } |
| 492 | |
| 493 | |
| 494 | def load(init_params: dict, init_file: str, model_cfg: dict, dont_load: Sequence[str] = ()) -> dict: |