DPointNet custom losses and callbacks#

Custom objectives and training callbacks are ordinary parts of an experiment. You do not need to edit BMTK to add them. This page shows the complete path: define a Python class, register its name, select it in JSON, and import the registration code before loading the model.

Start in the writable copy of examples/dpointnet_all2all used in DPointNet. The example below retains that network and its rand_inputs generator. See DPointNet training for the general training configuration.

Define and register the extensions#

Create custom_training.py in the example working directory with the following contents:

import json
import math
from pathlib import Path

import tensorflow as tf
from bmtk.simulator.dpointnet import add_loss_module
from bmtk.simulator.dpointnet.callbacks import Callbacks, add_callbacks


class MeanRateMSE:
    def __init__(self, rnn, target_hz, cost=1.0, **kwargs):
        self.dt = float(rnn.dt)
        self.target_hz = float(target_hz)
        self.cost = float(cost)
        if not math.isfinite(self.dt) or self.dt <= 0:
            raise ValueError("dt must be finite and positive")
        for name in ("target_hz", "cost"):
            value = getattr(self, name)
            if not math.isfinite(value) or value < 0:
                raise ValueError(f"{name} must be finite and nonnegative")

    def __call__(self, spikes, **kwargs):
        spikes = tf.cast(spikes, tf.float32)
        rates_hz = tf.reduce_mean(spikes, axis=(0, 1)) * (1000.0 / self.dt)
        return self.cost * tf.reduce_mean(tf.square(rates_hz - self.target_hz))


add_loss_module(MeanRateMSE, module_name="MeanRateMSE")


class StepTraceCallbacks(Callbacks):
    def __init__(self, rnn, run_label="trial", stop_below=None, **kwargs):
        super().__init__(rnn=rnn, **kwargs)
        self.run_label = run_label
        self.stop_below = None if stop_below is None else float(stop_below)
        if self.stop_below is not None and not math.isfinite(self.stop_below):
            raise ValueError("stop_below must be finite or None")
        self.trace_path = Path(self.callbacks_dir) / "step_events.jsonl"

    def on_train_begin(self):
        super().on_train_begin()
        self.trace_path.parent.mkdir(parents=True, exist_ok=True)
        self.trace_path.write_text("", encoding="utf-8")

    def on_step_end(self, loss_vals):
        super().on_step_end(loss_vals)
        record = {
            "run_label": self.run_label,
            "epoch": self.epoch_num,
            "step": self.epoch_step_num,
            "total_loss": float(loss_vals["__total_loss"]),
        }
        with self.trace_path.open("a", encoding="utf-8") as stream:
            stream.write(json.dumps(record, allow_nan=False) + "\n")

    def on_epoch_end(self, validation_losses):
        stop = super().on_epoch_end(validation_losses)
        score = float(validation_losses["__total_loss"])
        if not math.isfinite(score):
            raise ValueError("Non-finite validation total")
        return bool(stop) or (
            self.stop_below is not None and score < self.stop_below
        )


add_callbacks(StepTraceCallbacks, name="StepTraceCallbacks")

The custom loss fits each neuron’s mean firing rate, pooled across batch and time, to one scalar target in Hz. It is not identical to the built-in TargetFiringRate, which compares rates separately for each trial. The custom callback adds a JSON-lines trace and optional epoch-level stopping while preserving the standard loss tables and weight exports.

Loss interface#

The JSON loader constructs a loss as LossClass(rnn=network, **loss_entry). The constructor therefore needs rnn, the configurable arguments it uses, and **kwargs for loader fields such as module and enabled. Explicit registration above supplies the lookup name; no Keras loss subclass or module() method is required when module_name is given.

During training and validation, DPointNet calls the object with keyword arguments including spikes, voltages, model_state and y. Use __call__(self, spikes, **kwargs) for a spike-only loss. This is not Keras’s loss(y_true, y_pred) interface; y is input-generator metadata, not an automatically loaded target array.

  • spikes has shape [batch, time, neurons]. Output columns follow the model’s neuron ordering, not arbitrary SONATA node IDs.

  • Return one differentiable scalar tensor. The training engine adds the configured loss terms; coefficients such as cost belong to your loss.

  • Keep computation in TensorFlow. Converting spikes to NumPy or Python values inside the loss breaks the gradient path. The example casts to FP32 without detaching gradients, including when model computation uses mixed precision.

  • Load fixed target files and construct masks in __init__, then keep the loss calculation free of file I/O and Python-side history updates. Graph tracing, validation and checkpoint recomputation can invoke the loss outside the sequence of accepted optimizer updates.

  • For a loss needing neuron-resolved voltages, explicitly set rnn_cell_params.return_voltage_sequences=true. The default compact voltage statistic is not a full [batch, time, neurons] voltage tensor.

Use a unique registration name. Reusing a built-in name replaces its implementation in that Python process; that should be an intentional choice, not a way to select a different objective accidentally.

Callback interface#

DPointNet uses its own Callbacks interface, not tf.keras.callbacks.Callback. Subclass the DPointNet class and forward rnn and standard options to super().__init__. For hooks you override, call the parent hook to retain logging, counters, timings, saved weights and exports. Unchanged hooks are inherited automatically.

Training hooks#

Method

When it runs

on_train_begin()

Once before the first epoch.

on_epoch_start()

Before an epoch; the parent increments epoch_num.

on_step_start()

Before a training step; the parent increments step counters.

on_step_end(loss_vals)

After the training step, with its loss and metric dictionary.

on_epoch_end(validation_losses)

After evaluation; returning true stops the epoch loop.

on_train_end(metrics=None, normalizers=None)

At normal training completion; the parent exports saved weights.

Both loss dictionaries contain "__total_loss" and nested dictionaries such as loss_vals["custom_rates"]["mean_rates"] for named terms. Reserved metric keys begin with __. Convert scalar tensors to Python values for host-side logging, as the callback does, not inside the loss.

The example’s stop_below=None disables early stopping. A numeric threshold compares the reported validation total, including regularizers; it does not create a held-out split. Choose the stopping score for your experiment. Use a new callback output directory for a new run: this example resets its trace at on_train_begin and does not resume a previous training trajectory.

Select the extensions in configuration#

Save the following JSON fragment as custom_training_config.json. The launcher below merges its training fields into the complete config.train.all.json. Thus the original network, inputs, initial state, two epochs and three steps per epoch remain available.

{
  "training": {
    "training_approach": "single",
    "gradient_checkpointing": false,
    "learning_rate": 0.001,
    "callbacks": {
      "class": "StepTraceCallbacks",
      "callbacks_dir": "custom_training_outputs",
      "run_label": "rate_trial",
      "stop_below": null,
      "verbose": "on_epoch",
      "memory_report": "off",
      "epoch_store_weights": "latest",
      "sonata_output_dir": "trained_network",
      "losses_table_csv": "losses.csv"
    },
    "parameters": [
      {
        "name": "custom_rates",
        "inputs": ["rand_inputs"],
        "loss_functions": {
          "mean_rates": {
            "module": "MeanRateMSE",
            "target_hz": 15.0,
            "cost": 0.01
          },
          "voltage": {
            "module": "VoltageRegularization",
            "voltage_cost": 1.5,
            "penalty_mode": "range"
          }
        }
      }
    ]
  }
}

Here "mean_rates" is a reporting label, whereas "MeanRateMSE" is the registered implementation name. Losses belong in each condition’s loss_functions; callbacks belongs under training and selects one callback class. Combine additional hooks in that class rather than supplying a Keras callback list.

The coefficients are illustrative. Checkpointing is disabled only for this legacy example’s refractory-state dtype limitation, explained in DPointNet training; it is not a general requirement for custom losses.

Import registration code and run#

Save this as run_custom_training.py beside the other two files:

import json
import os
from pathlib import Path

os.environ.pop("TF_GPU_ALLOCATOR", None)

import custom_training
from bmtk.simulator import dpointnet


def main():
    config = dpointnet.Config.from_json("config.train.all.json")
    fragment = json.loads(Path("custom_training_config.json").read_text())
    config["training"].update(fragment["training"])
    config.build_env()
    rnn = dpointnet.RNN.from_config(config)
    try:
        rnn.run()
    finally:
        rnn.cleanup()


if __name__ == "__main__":
    main()

Run from the example directory, using the environment with your intended BMTK installation:

CUDA_VISIBLE_DEVICES="" python run_custom_training.py

Importing custom_training executes both registrations before RNN.from_config looks up the JSON names. JSON does not import an arbitrary Python file automatically. Launching the unchanged generic entrypoint in another process will not inherit these registrations.

With early stopping disabled, the example makes six optimizer updates and then runs inference. Check:

  • custom_training_outputs/step_events.jsonl: six records containing the run label, epoch, step and total loss.

  • custom_training_outputs/losses.csv: individual mean_rates and voltage terms as well as totals.

  • custom_training_outputs/trained_network/: exported trained edge files; see DPointNet training for reusing them with the original network assets.

For a missing loss-module or callback-class error, first check the spelling and whether the registration module was imported in the same Python process before model loading. For disconnected gradients, check that the loss returns a tensor and does not use NumPy, .numpy() or stop_gradient on the prediction path. Test the scalar value and gradients on a small known tensor before applying a new loss to a large network.