Converts a string like "B" or "B/32" into a params dict.
(variant)
| 814 | |
| 815 | |
| 816 | def decode_variant(variant): |
| 817 | """Converts a string like "B" or "B/32" into a params dict.""" |
| 818 | if variant is None: |
| 819 | return {} |
| 820 | |
| 821 | v, patch = variant, {} |
| 822 | if "/" in variant: |
| 823 | v, patch = variant.split("/") |
| 824 | patch = {"patch_size": (int(patch), int(patch))} |
| 825 | |
| 826 | return { |
| 827 | # pylint:disable=line-too-long |
| 828 | # Reference: Table 2 of https://arxiv.org/abs/2106.04560. |
| 829 | "width": {"mu": 32, "Ti": 192, "S": 384, "M": 512, "B": 768, "L": 1024, "So400m": 1152, "H": 1280, "g": 1408, "g-opt": 1536, "G": 1664, "G-opt": 1536, "e": 1792}[v], |
| 830 | "depth": {"mu": 1, "Ti": 12, "S": 12, "M": 12, "B": 12, "L": 24, "So400m": 27, "H": 32, "g": 40, "g-opt": 40, "G": 48, "G-opt": 48, "e": 56}[v], |
| 831 | "mlp_dim": {"mu": 128, "Ti": 768, "S": 1536, "M": 2048, "B": 3072, "L": 4096, "So400m": 4304, "H": 5120, "g": 6144, "g-opt": 6144, "G": 8192, "G-opt": 8192, "e": 15360}[v], |
| 832 | "num_heads": {"mu": 2, "Ti": 3, "S": 6, "M": 8, "B": 12, "L": 16, "So400m": 16, "H": 16, "g": 16, "g-opt": 16, "G": 16, "G-opt": 16, "e": 16}[v], |
| 833 | # pylint:enable=line-too-long |
| 834 | **patch |
| 835 | } |
| 836 | |
| 837 | |
| 838 | def resample_posemb(old, new): |