File size: 622 Bytes
3a5d3e5 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 | import yaml
import sys
from pathlib import Path
ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT / "model" / "scripts"))
import viu1_dense
with open(ROOT / "model" / "configs" / "viu1_dense_config.yaml", "r") as f:
cfg = yaml.safe_load(f)
args = viu1_dense.load_args(ROOT / "model" / "configs" / "viu1_dense_config.yaml")
model = viu1_dense.Viu1Dense(args)
total, active = model.param_counts()
print(f"Total params in ViuMini-Dense-360M: {total:,} ({total/1e6:.1f}M) | Active: {active:,}")
print(f"Layers: {args.n_layers}, Dim: {args.dim}, Heads: {args.n_heads}, KV Heads: {args.n_kv_heads}")
|