Converts a string like "B" or "B/32" into a params dict.
(variant)
| 582 | |
| 583 | |
| 584 | def decode_variant(variant): |
| 585 | """Converts a string like "B" or "B/32" into a params dict.""" |
| 586 | if variant is None: |
| 587 | return {} |
| 588 | |
| 589 | v = variant |
| 590 | |
| 591 | return { |
| 592 | # pylint:disable=line-too-long |
| 593 | # Reference: Table 2 of https://arxiv.org/abs/2106.04560. |
| 594 | # from text transformer |
| 595 | "width": {"Ti": 192, "S": 384, "M": 512, "B": 512, "L": 768, "So400m": 1152, "H": 1024, "g": 1024, "G": 1664, "e": 1792}[v], |
| 596 | "depth": {"Ti": 12, "S": 12, "M": 12, "B": 12, "L": 12, "So400m": 27,"H": 24, "g": 24, "G": 48, "e": 56}[v], |
| 597 | "mlp_dim": {"Ti": 768, "S": 1536, "M": 2048, "B": 2048, "L": 3072, "So400m": 4304,"H": 4096, "g": 4096, "G": 8192, "e": 15360}[v], |
| 598 | "num_heads": {"Ti": 3, "S": 6, "M": 8, "B": 8, "L": 12, "So400m": 16, "H": 16, "g": 16, "G": 16, "e": 16}[v], |
| 599 | # pylint:enable=line-too-long |
| 600 | } |
| 601 | |
| 602 | |
| 603 | def load(init_params, init_file, model_cfg, dont_load=()): # pylint: disable=invalid-name because we had to CamelCase above. |