Feature/dpointnet#461
Merged
Merged
Conversation
…eproduce old results
fixing bug in GLIF3Cell init, plus updating code so it will exactly r…
updating firingratestargets loss function
…inable background + loss core (#454) * Working version for inference. Input fix is also included. * Fixing OOM due to the scope and strategy * dpointnet: float16 training fixes, recurrent-weight shadow refresh, and config updates - glif3_cell.py: cast decay/current_factor to compute_dtype (float16 mixed precision), always keep a non-trainable recurrent-weight compute shadow + refresh_recurrent_weight_shadow() so trained weights propagate to the forward, and port calculate_input_current_from_firing_probabilities. - training.py: fix batched param slicing (advance pidx_beg), cast loss accumulation to float32, gradient-checkpointing (recompute_grad) plumbing, MirroredStrategy-compatible run. - configs: switch training configs to float16, migrate background input to poisson spikes (module: random), match reference loss weights (rate_cost 10000 / voltage 1.0 / sync 1.5), add missing tau_syns, point synaptic_models_dir at synaptic_models_updated. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> * dpointnet: enable input (background) weight training + per-loss core_radius Make input weight trainability follow the config and add a loss "core" option, matching the reference V1_GLIF_model training condition. - rnn_model.parse_input_mods_from_config: propagate each input's config 'trainable'/'weight_scale' onto the input network options, which the GLIF cell reads when creating the input weight Variable. Previously net_dict['options'] was overwritten so the config flag was a no-op and no input (incl. background) could be trained. This re-enables training the background "rest_of_brain" weights, matching the reference's train_noise. - loss_functions: add loss_utils.resolve_core_mask() and core_radius support to SpikeRateDistributionTarget, SynchronizationLoss and VoltageRegularization so the loss can be restricted to a central core (matches reference loss_core_radius). isolate_core_neurons now falls back to identity neuron order when the network has no 'tf_id_to_bmtk_id' (dpointnet keeps SONATA order). - example config.training.multispikes.batched.json: match reference condition (exp_adam, constant lr=0.005, dampening 0.1/0.1, core_radius 200, trainable background @250Hz). osi_dsi loss intentionally omitted. All code changes are additive/backward-compatible (core_radius defaults to None; input trainability acts only when the config sets it). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> * dpointnet: memory-efficient custom-gradient input current Give the input (LGN/background) synaptic current the same @tf.custom_gradient treatment as the recurrent current, so per-timestep activations are recomputed in backward instead of stacked across the whole sequence. - Add module-level calculate_input_currents (mirrors calculate_synaptic_currents): forward reads a compute-dtype shadow, backward recomputes the basis gather and returns the gradient w.r.t. the input weight master (supports multi-spike inputs via the n_pre_spikes factor; no gradient to the external input x_t). - Add a non-trainable compute-dtype shadow for trainable input weights and refresh it each step (extends refresh_recurrent_weight_shadow); avoids fp16 casts and reading a trainable MirroredVariable inside the RNN while_loop. - Route calculate_input_current_from_spikes through the new function. Verified: the custom gradient matches plain autodiff on the background weights (relative diff ~4e-4). Training still descends at lr=0.005. Peak memory (batch1, seq500, full losses): no-checkpointing 10.2 -> 9.75 GB; checkpointing path 13.7 -> 10.8 GB. Background-input activations no longer dominate at high firing rates; savings scale with batch size. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> * dpointnet: address PR review (poisson_spikes dtype, validation batched slicing) - PoissonSpikes.create_generator now yields spikes in the requested dtype so it matches the dataset output_signature; tf.data.from_generator requires the numpy dtype to match the TensorSpec dtype, which previously errored under mixed-float16. - _validation_step (batched): advance pidx_beg per parameter so each parameter's validation losses use its own batch segment (same fix already applied in _train_step_batched); previously parameter 2+ used the [0:pidx_end] slice, making per-parameter and total validation losses incorrect. (The review's note that validation bypasses gradient_checkpointing is not a memory issue: _validation_step is forward-only with no GradientTape, so no activations are retained.) Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> --------- Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
This reverts commit db54cfe.
* Update dpointnet training approach modes * Fix dpointnet voltage regularization gradients (added after approval) * Address dpointnet training mode review comments (also added after approval)
* Stabilize dpointnet paper training path * Stabilize dpointnet training diagnostics and export ordering * Updating the basis function treatment. * Correct poisson * Correcting the timing that EMA is calculated * Gray screen cache + better GPU output * including LGN generation in tf.function * Finally fixing memory retention bug * Defaulting the BFC memory allocator. Giving up malloc * Implementing annulus loss, and some other changes to match with the reference code. * Config for reproducing old results * Making ccw default * Fix LGN tensor path and OSI JIT stability Keep drifting-grating random orientation on the TensorFlow tensor path to avoid Python scalar extraction in LGN generation. Stabilize OSI/DSI XLA computation with dynamic batch sizing and separate compiled core/annulus paths. Cast bool checkpoint inputs to the model dtype so recompute_grad can run, while leaving checkpointing opt-in. * Recover transient input-state generation failures Retry grey-screen initial-state generation on transient TensorFlow shape errors by rebuilding its input iterator. If retries are exhausted after at least one successful state, reuse the last generated state instead of aborting a long run at an epoch boundary. * Avoid repeated input-state recovery failures After a transient input-state generation error is successfully recovered, reuse the recovered state for later epochs instead of probing the failing path every epoch. * Revert "Avoid repeated input-state recovery failures" This reverts commit e969d4c. * Use vector shape for scalar LGN RNG Use shape=[] for scalar theta and phase random draws so stateless RNG receives a vector shape argument while preserving scalar Tensor outputs. * Run LGN seed fold-in on CPU Keep TensorFlow stateless_fold_in semantics but place the small seed-mixing op on CPU to avoid intermittent GPU/BFC StatelessRandomUniformFullIntV2 shape errors from tf.data generators. * updating paper config * Simplify dpointnet v1 example setup Keep one training config and one verified 45-degree 3s inference config at the example top level. Move older tracked configs and duplicate train/inference runners under misc_setup for later cleanup. * Rename background spike input modules Expose the internally generated Poisson drive as poisson_spikes_internal and the external Bernoulli generator as bernoulli_spikes. Keep noisy_current and random as compatibility aliases while updating canonical and all2all configs to the clearer names. * making cuda_malloc_async default * Restore dpointnet input prefetch Restore tf.data prefetching in DataIterator after A100 testing showed it improves full seq_len=500 training step time while keeping the run stable. Also remove the stale initial_stateX block from the canonical paper training config. * Sync loss to ms * Fixing a bug that resets noise every epoch
Enabling multi-parameter inference and fixed phase lgn runs
updating unit tests to include tf
…/bmtk into feature/dpointnet
Feature/dpointnet
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
No description provided.