| """Numpy implementation of UL-UNAS DPGRNN (grouped dual-path RNN). |
| |
| Pulsar2 cannot lower the large-batch GRU transposes in this block to AX650 |
| (``AxTranspose: size 16 cannot explode into 256``), so the two DPGRNN blocks |
| run on the host between the NPU encoder and decoder. They are tiny: 2 blocks, |
| 6 GRU directions, ~60 k parameters, < 5 ms for a 4.8 s chunk. |
| """ |
| import numpy as np |
|
|
|
|
| def _sigmoid(x): |
| return 1.0 / (1.0 + np.exp(-x)) |
|
|
|
|
| def _gru(x, p, prefix, reverse=False): |
| w_ih = p[prefix + ".weight_ih_l0" + ("_reverse" if reverse else "")] |
| w_hh = p[prefix + ".weight_hh_l0" + ("_reverse" if reverse else "")] |
| b_ih = p[prefix + ".bias_ih_l0" + ("_reverse" if reverse else "")] |
| b_hh = p[prefix + ".bias_hh_l0" + ("_reverse" if reverse else "")] |
| if reverse: |
| x = x[:, ::-1] |
| b, t, _ = x.shape |
| hsize = w_hh.shape[1] |
| gx = x @ w_ih.T + b_ih |
| h = np.zeros((b, hsize), dtype=np.float32) |
| out = np.empty((b, t, hsize), dtype=np.float32) |
| for i in range(t): |
| gh = h @ w_hh.T + b_hh |
| r = _sigmoid(gx[:, i, :hsize] + gh[:, :hsize]) |
| z = _sigmoid(gx[:, i, hsize:2 * hsize] + gh[:, hsize:2 * hsize]) |
| n = np.tanh(gx[:, i, 2 * hsize:] + r * gh[:, 2 * hsize:]) |
| h = (1.0 - z) * n + z * h |
| out[:, i] = h |
| if reverse: |
| out = out[:, ::-1] |
| return out |
|
|
|
|
| def _grnn(x, p, prefix): |
| x1, x2 = np.split(x, 2, axis=-1) |
| y1 = _gru(x1, p, prefix + ".rnn1") |
| if (prefix + ".rnn1.weight_ih_l0_reverse") in p: |
| y1b = _gru(x1, p, prefix + ".rnn1", reverse=True) |
| y1 = np.concatenate([y1, y1b], axis=-1) |
| y2 = _gru(x2, p, prefix + ".rnn2") |
| if (prefix + ".rnn2.weight_ih_l0_reverse") in p: |
| y2b = _gru(x2, p, prefix + ".rnn2", reverse=True) |
| y2 = np.concatenate([y2, y2b], axis=-1) |
| return np.concatenate([y1, y2], axis=-1) |
|
|
|
|
| def _layer_norm(x, w, b): |
| mean = x.mean(axis=(-2, -1), keepdims=True) |
| var = x.var(axis=(-2, -1), keepdims=True) |
| return (x - mean) / np.sqrt(var + 1e-8) * w + b |
|
|
|
|
| def _linear(x, p, prefix): |
| return x @ p[prefix + ".weight"].T + p[prefix + ".bias"] |
|
|
|
|
| class DPGRNN: |
| def __init__(self, npz_path, blocks=2): |
| self.p = dict(np.load(npz_path)) |
| self.blocks = blocks |
|
|
| def _one(self, x, i): |
| p = self.p |
| pre = f"b{i}." |
| x = x.transpose(0, 2, 3, 1) |
| b, t, f, c = x.shape |
| intra = _grnn(x.reshape(b * t, f, c), p, pre + "intra_rnn") |
| intra = _linear(intra, p, pre + "intra_fc").reshape(b, t, f, c) |
| intra = _layer_norm(intra, p[pre + "intra_ln.weight"], p[pre + "intra_ln.bias"]) |
| intra_out = x + intra |
|
|
| x = intra_out.transpose(0, 2, 1, 3) |
| inter = _grnn(x.reshape(b * f, t, c), p, pre + "inter_rnn") |
| inter = _linear(inter, p, pre + "inter_fc").reshape(b, f, t, c) |
| inter = inter.transpose(0, 2, 1, 3) |
| inter = _layer_norm(inter, p[pre + "inter_ln.weight"], p[pre + "inter_ln.bias"]) |
| out = intra_out + inter |
| return out.transpose(0, 3, 1, 2) |
|
|
| def __call__(self, x): |
| for i in range(self.blocks): |
| x = self._one(x, i) |
| return x |
|
|
|
|
| def extract_params(model, out_path): |
| """Extract DPGRNN weights from a torch ULUNAS model into a small .npz.""" |
| sd = model.state_dict() |
| out = {} |
| for i, block in enumerate(model.dpgrnn): |
| pre = f"b{i}." |
| for part in ("intra_rnn", "inter_rnn"): |
| for rnn in ("rnn1", "rnn2"): |
| p = f"dpgrnn.{i}.{part}.{rnn}" |
| for suffix in ("weight_ih_l0", "weight_hh_l0", "bias_ih_l0", "bias_hh_l0"): |
| key = f"{p}.{suffix}" |
| if key in sd: |
| out[pre + f"{part}.{rnn}.{suffix}"] = sd[key].detach().numpy().astype(np.float32) |
| if key + "_reverse" in sd: |
| out[pre + f"{part}.{rnn}.{suffix}_reverse"] = sd[key + "_reverse"].detach().numpy().astype(np.float32) |
| for part in ("intra_fc", "inter_fc"): |
| p = f"dpgrnn.{i}.{part}" |
| out[pre + f"{part}.weight"] = sd[p + ".weight"].detach().numpy().astype(np.float32) |
| out[pre + f"{part}.bias"] = sd[p + ".bias"].detach().numpy().astype(np.float32) |
| for part in ("intra_ln", "inter_ln"): |
| p = f"dpgrnn.{i}.{part}" |
| out[pre + f"{part}.weight"] = sd[p + ".weight"].detach().numpy().astype(np.float32) |
| out[pre + f"{part}.bias"] = sd[p + ".bias"].detach().numpy().astype(np.float32) |
| np.savez(out_path, **out) |
| return len(out) |
|
|
