diff options
| author | tslil <tslil@posteo.de> | 2026-09-19 20:24:08 +0100 |
|---|---|---|
| committer | tslil <tslil@posteo.de> | 2026-09-19 20:32:46 +0100 |
| commit | 9f5f953a95c8617cd2fdd24f7402af632c1b07c2 (patch) | |
| tree | d78abd447065baf3d100935a2cfbb181d50b0fb6 /src/nn_train.py | |
| parent | fe7268110d93e2abe4d4b51ac64e83521a8d6382 (diff) | |
transposing the weights gives ~2.5x speedup
Diffstat (limited to 'src/nn_train.py')
| -rwxr-xr-x | src/nn_train.py | 11 |
1 files changed, 8 insertions, 3 deletions
diff --git a/src/nn_train.py b/src/nn_train.py index f82a369..23cf38b 100755 --- a/src/nn_train.py +++ b/src/nn_train.py @@ -486,9 +486,14 @@ def fast_tanh(x: Tensor) -> Tensor: x2 = x * x poly = x * ( 0.997136032 - + x2 * ( + + x2 + * ( -0.305215124 - + x2 * (0.0832466371 + x2 * (-0.0139947924 + x2 * (0.00123106314 + x2 * -4.29748644e-05))) + + x2 + * ( + 0.0832466371 + + x2 * (-0.0139947924 + x2 * (0.00123106314 + x2 * -4.29748644e-05)) + ) ) ) return torch.where(x > 3.0, 1.0, torch.where(x < -3.0, -1.0, poly)) @@ -543,7 +548,7 @@ class Net(Module): f"#ifndef {guard}\n#define {guard}\n", f'#include "weights_{size}.h"\n\n', f"const float {prefix}_embedding[{len(E)}] =\n{{\n{_fmt_flat(E, per_line=7)}\n}};\n\n", - f"const float {prefix}_dense_weights[{upper}_DENSE_NUM][{upper}_INP_NUM] NN_ALIGNED =\n{{\n{_fmt_matrix(W)}\n}};\n\n", + f"const float {prefix}_dense_weights[{upper}_INP_NUM][{upper}_DENSE_NUM] NN_ALIGNED =\n{{\n{_fmt_matrix(W.T)}\n}};\n\n", f"const float {prefix}_dense_biases[{upper}_DENSE_NUM] =\n{{\n{_fmt_flat(B)}\n}};\n\n", f"const float {prefix}_logit_weights[{upper}_DENSE_NUM] NN_ALIGNED =\n{{\n{_fmt_flat(LW)}\n}};\n\n", f"const float {prefix}_logit_bias = {LB:.9g};\n", |