Simple-STNDT使用Transformer进行Spike信号的表征学习（一）数据处理篇

本文主要是介绍Simple-STNDT使用Transformer进行Spike信号的表征学习（一）数据处理篇，希望对大家解决编程问题提供一定的参考价值，需要的开发者们随着小编来一起学习吧！

文章目录

- 1.数据处理部分
- - 1.1 下载数据集
  - 1.2 数据集预处理
  - 1.3 划分train-val并创建Dataset对象
  - 1.4 掩码mask操作

数据、评估标准见NLB2021
https://neurallatents.github.io/

以下代码依据
https://github.com/trungle93/STNDT

原代码使用了 Ray+Config文件进行了参数搜索，库依赖较多，数据流过程不明显，代码冗杂，这里进行了抽丝剥茧，将其中最核心的部分提取出来。

1.数据处理部分

1.1 下载数据集

需要依赖 pip install dandi
downald.py

root = "D:/NeuralLatent/"
def downald_data():from dandi.download import downloaddownload("https://dandiarchive.org/dandiset/000128", root)download("https://dandiarchive.org/dandiset/000138", root)download("https://dandiarchive.org/dandiset/000139", root)download("https://dandiarchive.org/dandiset/000140", root)download("https://dandiarchive.org/dandiset/000129", root)download("https://dandiarchive.org/dandiset/000127", root)download("https://dandiarchive.org/dandiset/000130", root)

1.2 数据集预处理

需要依赖官方工具包pip install nlb_tools
主要是加载锋值序列数据，将其采样为5ms的时间槽
preprocess.py

## 以下为参数示例
# data_path = root + "/000129/sub-Indy/"
# dataset_name = "mc_rtt"
## 注意 "./data" 必须提前创建好from nlb_tools.make_tensors import make_train_input_tensors, make_eval_input_tensors, combine_h5def preprocess(data_path, dataset_name=None):dataset = NWBDataset(datapath)bin_width = 5dataset.resample(bin_width)make_train_input_tensors(dataset, dataset_name=dataset_name, trial_split="train", include_behavior=True, include_forward_pred=True, save_file=True,save_path=f"./data/{dataset_name}_train.h5")make_eval_input_tensors(dataset, dataset_name=dataset_name, trial_split="val", save_file=True, save_path=f"./data/{dataset_name}_val.h5")combine_h5([f"./data/{dataset_name}_train.h5", f"./data/{dataset_name}_val.h5"], save_path=f"./data/{dataset_name}_full.h5")## './data/mc_rtt_full.h5' 将成为后续的主要分析数据

1.3 划分train-val并创建Dataset对象

读取'./data/mc_rtt_full.h5'中的数据并创建dataset
dataset.py

import h5py
import numpy as np
import torch
from torch.utils import data
# data_path = "./data/mc_rtt_full.h5"class SpikesDataset(data.Dataset):def __init__(self, spikes, heldout_spikes, forward_spikes) -> None:self.spikes = spikesself.heldout_spikes = heldout_spikesself.forward_spikes = forward_spikesdef __len__(self):return self.spikes.size(0)def __getitem__(self, index):r"""Return spikes and rates, shaped T x N (num_neurons)"""return self.spikes[index], self.heldout_spikes[index], self.forward_spikes[index]def make_datasets(data_path):with h5py.File(data_path, 'r') as h5file:h5dict = {key: h5file[key][()] for key in h5file.keys()}if 'eval_spikes_heldin' in h5dict: # NLB dataget_key = lambda key: h5dict[key].astype(np.float32)train_data = get_key('train_spikes_heldin')train_data_fp = get_key('train_spikes_heldin_forward')train_data_heldout_fp = get_key('train_spikes_heldout_forward')train_data_all_fp = np.concatenate([train_data_fp, train_data_heldout_fp], -1)valid_data = get_key('eval_spikes_heldin')train_data_heldout = get_key('train_spikes_heldout')if 'eval_spikes_heldout' in h5dict:valid_data_heldout = get_key('eval_spikes_heldout')else:valid_data_heldout = np.zeros((valid_data.shape[0], valid_data.shape[1], train_data_heldout.shape[2]), dtype=np.float32)if 'eval_spikes_heldin_forward' in h5dict:valid_data_fp = get_key('eval_spikes_heldin_forward')valid_data_heldout_fp = get_key('eval_spikes_heldout_forward')valid_data_all_fp = np.concatenate([valid_data_fp, valid_data_heldout_fp], -1)else:valid_data_all_fp = np.zeros((valid_data.shape[0], train_data_fp.shape[1], valid_data.shape[2] + valid_data_heldout.shape[2]), dtype=np.float32)train_dataset = SpikesDataset(torch.tensor(train_data).long(),            # [810, 120, 98]torch.tensor(train_data_heldout).long(),    # [810, 120, 32]torch.tensor(train_data_all_fp).long(),     # [810, 40, 130])val_dataset = SpikesDataset(torch.tensor(valid_data).long(),            # [810, 120, 98]torch.tensor(valid_data_heldout).long(),    # [810, 120, 32]torch.tensor(valid_data_all_fp).long(),     # [810, 40, 130])return train_dataset, val_dataset

1.4 掩码mask操作

dataset.py

# Some infeasibly high spike count
UNMASKED_LABEL = -100def mask_batch(batch, heldout_spikes, forward_spikes):batch = batch.clone() # make sure we don't corrupt the input data (which is stored in memory)mask_ratio = 0.31254mask_random_ratio = 0.876mask_token_ratio = 0.527labels = batch.clone()mask_probs = torch.full(labels.shape, mask_ratio)# If we want any tokens to not get masked, do it here (but we don't currently have any)mask = torch.bernoulli(mask_probs)mask = mask.bool()labels[~mask] = UNMASKED_LABEL  # No ground truth for unmasked - use this to mask loss# We use random assignment so the model learns embeddings for non-mask tokens, and must rely on context# Most times, we replace tokens with MASK tokenindices_replaced = torch.bernoulli(torch.full(labels.shape, mask_token_ratio)).bool() & maskbatch[indices_replaced] = 0# Random % of the time, we replace masked input tokens with random value (the rest are left intact)indices_random = torch.bernoulli(torch.full(labels.shape, mask_random_ratio)).bool() & mask & ~indices_replacedrandom_spikes = torch.randint(batch.max(), labels.shape, dtype=torch.long)batch[indices_random] = random_spikes[indices_random]# heldout spikes are all maskedbatch = torch.cat([batch, torch.zeros_like(heldout_spikes)], -1)labels = torch.cat([labels, heldout_spikes.to(batch.device)], -1)batch = torch.cat([batch, torch.zeros_like(forward_spikes)], 1)labels = torch.cat([labels, forward_spikes.to(batch.device)], 1)# Leave the other 10% alonereturn batch, labels