feat(wasm): add batch inference, extended training, pin mask, eval loss, and layer stats bindings
Five new C functions for the WASM module: - nisps_mlp_infer_batch: N-point batch inference in a single call - nisps_mlp_train_ex: training with per-iteration loss history output - nisps_mlp_move_weights_ex: moveWeights with output pin mask to skip pinned nodes - nisps_mlp_eval_loss: compute MSE loss without updating weights - nisps_mlp_get_layer_stats: per-layer weight magnitude, dead, and saturation stats
This commit is contained in:
parent
44fc974425
commit
f225c7c2a6
4 changed files with 188 additions and 1 deletions
|
|
@ -22,6 +22,11 @@ emcc "$SCRIPT_DIR/nisps_bindings.cpp" \
|
||||||
"_nisps_mlp_train",
|
"_nisps_mlp_train",
|
||||||
"_nisps_mlp_draw_weights_spread",
|
"_nisps_mlp_draw_weights_spread",
|
||||||
"_nisps_mlp_move_weights_spread",
|
"_nisps_mlp_move_weights_spread",
|
||||||
|
"_nisps_mlp_infer_batch",
|
||||||
|
"_nisps_mlp_train_ex",
|
||||||
|
"_nisps_mlp_move_weights_ex",
|
||||||
|
"_nisps_mlp_eval_loss",
|
||||||
|
"_nisps_mlp_get_layer_stats",
|
||||||
"_nisps_alloc",
|
"_nisps_alloc",
|
||||||
"_nisps_free",
|
"_nisps_free",
|
||||||
"_nisps_alloc_int",
|
"_nisps_alloc_int",
|
||||||
|
|
|
||||||
File diff suppressed because one or more lines are too long
Binary file not shown.
|
|
@ -6,6 +6,12 @@
|
||||||
#include <cstdlib>
|
#include <cstdlib>
|
||||||
#include <cmath>
|
#include <cmath>
|
||||||
|
|
||||||
|
// Helper subclass to access protected members for extended training
|
||||||
|
struct NispsMLPAccessor : public nisps::MLP<float> {
|
||||||
|
using nisps::MLP<float>::loss_fn_;
|
||||||
|
using nisps::MLP<float>::UpdateWeights;
|
||||||
|
};
|
||||||
|
|
||||||
extern "C" {
|
extern "C" {
|
||||||
|
|
||||||
// ---- Lifecycle ----
|
// ---- Lifecycle ----
|
||||||
|
|
@ -163,6 +169,182 @@ void nisps_mlp_move_weights_spread(void* ptr, float speed, float spread) {
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// ---- Batch inference ----
|
||||||
|
|
||||||
|
EMSCRIPTEN_KEEPALIVE
|
||||||
|
void nisps_mlp_infer_batch(void* ptr, float* inputs_flat, int n_points, int input_dim, float* outputs_flat, int output_dim) {
|
||||||
|
auto* mlp = static_cast<nisps::MLP<float>*>(ptr);
|
||||||
|
std::vector<float> in_vec(input_dim);
|
||||||
|
std::vector<float> out_vec;
|
||||||
|
for (int i = 0; i < n_points; i++) {
|
||||||
|
in_vec.assign(inputs_flat + i * input_dim, inputs_flat + (i + 1) * input_dim);
|
||||||
|
out_vec.clear();
|
||||||
|
mlp->GetOutput(in_vec, &out_vec, nullptr, true);
|
||||||
|
int n = output_dim < (int)out_vec.size() ? output_dim : (int)out_vec.size();
|
||||||
|
for (int j = 0; j < n; j++) {
|
||||||
|
outputs_flat[i * output_dim + j] = out_vec[j];
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// ---- Extended training with per-iteration loss history ----
|
||||||
|
|
||||||
|
EMSCRIPTEN_KEEPALIVE
|
||||||
|
int nisps_mlp_train_ex(void* ptr,
|
||||||
|
float* features_flat, int n_samples, int feature_dim,
|
||||||
|
float* labels_flat, int label_dim,
|
||||||
|
float* sample_weights,
|
||||||
|
float learning_rate, int max_iterations, float min_error,
|
||||||
|
float* loss_history_out) {
|
||||||
|
|
||||||
|
auto* mlp = static_cast<NispsMLPAccessor*>(static_cast<nisps::MLP<float>*>(ptr));
|
||||||
|
|
||||||
|
std::vector<std::vector<float>> features(n_samples);
|
||||||
|
std::vector<std::vector<float>> labels(n_samples);
|
||||||
|
|
||||||
|
for (int i = 0; i < n_samples; i++) {
|
||||||
|
features[i].assign(features_flat + i * feature_dim,
|
||||||
|
features_flat + (i + 1) * feature_dim);
|
||||||
|
labels[i].assign(labels_flat + i * label_dim,
|
||||||
|
labels_flat + (i + 1) * label_dim);
|
||||||
|
}
|
||||||
|
|
||||||
|
float sample_size_recip = 1.0f / n_samples;
|
||||||
|
|
||||||
|
int iter = 0;
|
||||||
|
for (iter = 0; iter < max_iterations; iter++) {
|
||||||
|
float iteration_loss = 0.0f;
|
||||||
|
|
||||||
|
for (int s = 0; s < n_samples; s++) {
|
||||||
|
float w = sample_weights ? sample_weights[s] : sample_size_recip;
|
||||||
|
|
||||||
|
std::vector<float> predicted_output;
|
||||||
|
std::vector<std::vector<float>> all_layers_activations;
|
||||||
|
|
||||||
|
mlp->GetOutput(features[s], &predicted_output, &all_layers_activations, false);
|
||||||
|
|
||||||
|
std::vector<float> deriv_error_output(predicted_output.size());
|
||||||
|
float loss = mlp->loss_fn_(labels[s], predicted_output, deriv_error_output, w);
|
||||||
|
|
||||||
|
iteration_loss += loss;
|
||||||
|
|
||||||
|
mlp->UpdateWeights(all_layers_activations, deriv_error_output, learning_rate);
|
||||||
|
}
|
||||||
|
|
||||||
|
if (!sample_weights) {
|
||||||
|
iteration_loss *= sample_size_recip;
|
||||||
|
}
|
||||||
|
|
||||||
|
loss_history_out[iter] = iteration_loss;
|
||||||
|
|
||||||
|
if (iteration_loss < min_error) {
|
||||||
|
iter++;
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
return iter;
|
||||||
|
}
|
||||||
|
|
||||||
|
// ---- moveWeights with output pin mask ----
|
||||||
|
|
||||||
|
EMSCRIPTEN_KEEPALIVE
|
||||||
|
void nisps_mlp_move_weights_ex(void* ptr, float speed, float spread, int* pin_mask, int n_outputs) {
|
||||||
|
auto* mlp = static_cast<nisps::MLP<float>*>(ptr);
|
||||||
|
float decay = 1.0f - 0.1f * spread;
|
||||||
|
size_t n_layers = mlp->m_layers.size();
|
||||||
|
for (size_t l = 0; l < n_layers; l++) {
|
||||||
|
int fan_in = mlp->m_layers[l].GetInputSize();
|
||||||
|
float xavier_scale = 1.0f / std::sqrt((float)fan_in);
|
||||||
|
float layer_scale = 1.0f * (1.0f - spread) + xavier_scale * spread;
|
||||||
|
bool is_output_layer = (l == n_layers - 1);
|
||||||
|
int node_idx = 0;
|
||||||
|
for (auto& node : mlp->m_layers[l].GetNodesChangeable()) {
|
||||||
|
// Skip pinned output nodes
|
||||||
|
if (is_output_layer && pin_mask && node_idx < n_outputs && pin_mask[node_idx] == 1) {
|
||||||
|
node_idx++;
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
for (size_t j = 0; j < node.m_weights.size(); j++) {
|
||||||
|
node.m_weights[j] *= decay;
|
||||||
|
float accum = 0;
|
||||||
|
for (int n = 0; n < 3; n++) {
|
||||||
|
accum += (float)rand() / RAND_MAX * 2.0f - 1.0f;
|
||||||
|
}
|
||||||
|
node.m_weights[j] += 3.0f * accum * speed * layer_scale;
|
||||||
|
}
|
||||||
|
node_idx++;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// ---- Evaluate loss without updating weights ----
|
||||||
|
|
||||||
|
EMSCRIPTEN_KEEPALIVE
|
||||||
|
float nisps_mlp_eval_loss(void* ptr,
|
||||||
|
float* features_flat, int n_samples, int feature_dim,
|
||||||
|
float* labels_flat, int label_dim,
|
||||||
|
float* sample_weights) {
|
||||||
|
|
||||||
|
auto* mlp = static_cast<NispsMLPAccessor*>(static_cast<nisps::MLP<float>*>(ptr));
|
||||||
|
float total_loss = 0.0f;
|
||||||
|
float sample_size_recip = 1.0f / n_samples;
|
||||||
|
|
||||||
|
for (int s = 0; s < n_samples; s++) {
|
||||||
|
float w = sample_weights ? sample_weights[s] : sample_size_recip;
|
||||||
|
|
||||||
|
std::vector<float> in_vec(features_flat + s * feature_dim,
|
||||||
|
features_flat + (s + 1) * feature_dim);
|
||||||
|
std::vector<float> label_vec(labels_flat + s * label_dim,
|
||||||
|
labels_flat + (s + 1) * label_dim);
|
||||||
|
|
||||||
|
std::vector<float> predicted_output;
|
||||||
|
mlp->GetOutput(in_vec, &predicted_output, nullptr, true);
|
||||||
|
|
||||||
|
std::vector<float> deriv_error_output(predicted_output.size());
|
||||||
|
float loss = mlp->loss_fn_(label_vec, predicted_output, deriv_error_output, w);
|
||||||
|
total_loss += loss;
|
||||||
|
}
|
||||||
|
|
||||||
|
if (!sample_weights) {
|
||||||
|
total_loss *= sample_size_recip;
|
||||||
|
}
|
||||||
|
|
||||||
|
return total_loss;
|
||||||
|
}
|
||||||
|
|
||||||
|
// ---- Per-layer weight statistics ----
|
||||||
|
|
||||||
|
EMSCRIPTEN_KEEPALIVE
|
||||||
|
void nisps_mlp_get_layer_stats(void* ptr, float* stats_out, int n_layers) {
|
||||||
|
auto* mlp = static_cast<nisps::MLP<float>*>(ptr);
|
||||||
|
int layers_to_process = n_layers < (int)mlp->m_layers.size() ? n_layers : (int)mlp->m_layers.size();
|
||||||
|
for (int l = 0; l < layers_to_process; l++) {
|
||||||
|
float sum_abs = 0.0f;
|
||||||
|
float max_abs = 0.0f;
|
||||||
|
int dead_count = 0;
|
||||||
|
int saturating_count = 0;
|
||||||
|
int total_weights = 0;
|
||||||
|
|
||||||
|
for (auto& node : mlp->m_layers[l].m_nodes) {
|
||||||
|
for (size_t j = 0; j < node.m_weights.size(); j++) {
|
||||||
|
float aw = std::fabs(node.m_weights[j]);
|
||||||
|
sum_abs += aw;
|
||||||
|
if (aw > max_abs) max_abs = aw;
|
||||||
|
if (aw < 0.01f) dead_count++;
|
||||||
|
if (aw > 3.0f) saturating_count++;
|
||||||
|
total_weights++;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
float inv_total = total_weights > 0 ? 1.0f / total_weights : 0.0f;
|
||||||
|
stats_out[l * 4 + 0] = sum_abs * inv_total; // mean absolute weight
|
||||||
|
stats_out[l * 4 + 1] = max_abs; // max absolute weight
|
||||||
|
stats_out[l * 4 + 2] = dead_count * inv_total; // fraction dead
|
||||||
|
stats_out[l * 4 + 3] = saturating_count * inv_total; // fraction saturating
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
// ---- Memory helpers ----
|
// ---- Memory helpers ----
|
||||||
|
|
||||||
EMSCRIPTEN_KEEPALIVE
|
EMSCRIPTEN_KEEPALIVE
|
||||||
|
|
|
||||||
Loading…
Reference in a new issue