Neural Networks Loading and Retraining

Introduction

YASS relies on two neural networks (NNs) for detection and denoising. The default config.yaml file relies on NNs previously trained on a specific hexagonal layout containing 7 channels of data (centre channel + 6 neighbouring channels in a hexagonal layout) and an acquisition rate of 20,000Hz. The default NNs can be found in the yass assets directories.

yass/src/yass/assets/nn_models/detect.pt

yass/src/yass/assets/nn_models/denoise.pt

Other Available Neural Networks

The default YASS installation comes with additional neural networks trained on two layouts: neuropixels (385 channels, 30Khz, 3ms spike width, 40um local neighbourhood) and UtahArray (30KHz, 3ms spike width and 700um neighbourhood). These networks can be added to the config.yaml file from these locations:

yass/assets/nn_models/NeuroPixels_30000Hz_3ms_40um/

yass/assets/nn_models/UtahArray_30000Hz_3ms_700um/

Retraining Neural Networks

We found that the default NNs are sufficient for most datasets that have an approximately hexagonal pattern. But in many cases the default NNs will not run on novel layouts due to the sampling rate or layout differences. Additionally, retraining the NNs on novel data and neuron templates could yield additional neurons and an improvement in unit quality.

We provide an automated method for retraining the NNs on novel data and geometries. During retraining YASS will perform a number of basic clustering steps including: "threshold-based" event detection and clustering of the largest amplitude neurons in your dataset. It will then use these neuron templates in a model to generate sufficient data to train the NNs. The YASS NNs can be retrained using the following command:

yass train config.yaml

YASS will generate new NN files and will automatically continue running the entire dataset using the updated NNs. The location of the retrained NNs is here:

tmp/nn_train/

Newly generated NNs can be used in any other dataset, by entering the full path of the NNs into the config.yaml file:

`detect:`

 `# model name, it can be an absolute path to a model `
 `# (e.g. /path/to/my/model.ckpt) or a name of one of models included in yass `
 `# (e.g. detect.pt). Check /src/yass/assets/nn_models for the list of available`
 `# models.`
 `filename: detect.pt  `

`denoise:`

 `# model name, same rule as detect filename`
 `filename: denoise.pt`