| import json | |
| import torch | |
| import torch.nn as nn | |
| class SmallAudioClassifierMLP(nn.Module): | |
| def __init__(self, input_dim=6144, num_labels=10, hidden_dims=(512, 256), dropout=0.2, activation="gelu"): | |
| super().__init__() | |
| act = {"relu": nn.ReLU, "gelu": nn.GELU, "silu": nn.SiLU}[activation] | |
| dims = [input_dim] + list(hidden_dims) | |
| layers = [] | |
| for i in range(len(dims) - 1): | |
| layers += [ | |
| nn.Linear(dims[i], dims[i+1]), | |
| nn.LayerNorm(dims[i+1]), | |
| act(), | |
| nn.Dropout(dropout), | |
| ] | |
| self.mlp = nn.Sequential(*layers) | |
| self.classifier = nn.Linear(dims[-1], num_labels) | |
| def forward(self, x): | |
| h = self.mlp(x) | |
| return self.classifier(h) | |
| def load_pretrained(model_dir: str, map_location="cpu"): | |
| with open(f"{model_dir}/config.json", "r", encoding="utf-8") as f: | |
| cfg = json.load(f) | |
| m = SmallAudioClassifierMLP( | |
| input_dim=cfg["input_dim"], | |
| num_labels=cfg["num_labels"], | |
| hidden_dims=tuple(cfg["hidden_dims"]), | |
| dropout=cfg["dropout"], | |
| activation=cfg["activation"], | |
| ) | |
| sd = torch.load(f"{model_dir}/pytorch_model.bin", map_location=map_location) | |
| m.load_state_dict(sd) | |
| m.eval() | |
| return m, cfg | |