parallelizing Neuron Simulator model

[Hananel-Hazan]

Hi,

I am running a training script similar to the HH example, but with the Neuron Simulator.
Running the script using num_workers=1 the script run fine. But, when using num_workers>1 I am getting the following error message:

Running 100 simulations in 100 batches.:   0%|          | 0/100 [00:00<?, ?it/s]joblib.externals.loky.process_executor._RemoteTraceback: 
"""
Traceback (most recent call last):
  File "/home/hananel/.local/lib/python3.8/site-packages/joblib/externals/loky/backend/queues.py", line 150, in _feed
    obj_ = dumps(obj, reducers=reducers)
  File "/home/hananel/.local/lib/python3.8/site-packages/joblib/externals/loky/backend/reduction.py", line 247, in dumps
    dump(obj, buf, reducers=reducers, protocol=protocol)
  File "/home/hananel/.local/lib/python3.8/site-packages/joblib/externals/loky/backend/reduction.py", line 240, in dump
    _LokyPickler(file, reducers=reducers, protocol=protocol).dump(obj)
  File "/home/hananel/.local/lib/python3.8/site-packages/joblib/externals/cloudpickle/cloudpickle_fast.py", line 548, in dump
    return Pickler.dump(self, obj)
TypeError: HocObject: Only Vector instance can be pickled
"""

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "/home/hananel/git repo/Neuron_Test/R859C/neuron-sbi.py", line 344, in <module>
    data[m + '-posterior'] = infer(test_parameters, prior, method=m, num_simulations=num_simulations, num_workers=2)
  File "/home/hananel/.local/lib/python3.8/site-packages/sbi/inference/base.py", line 72, in infer
    posterior = infer_(num_simulations=num_simulations)
  File "/home/hananel/.local/lib/python3.8/site-packages/sbi/inference/snpe/snpe_c.py", line 206, in __call__
    return super().__call__(**kwargs)
  File "/home/hananel/.local/lib/python3.8/site-packages/sbi/inference/snpe/snpe_base.py", line 172, in __call__
    theta, x = self._run_simulations(proposal, num_simulations)
  File "/home/hananel/.local/lib/python3.8/site-packages/sbi/inference/base.py", line 261, in _run_simulations
    x = self._batched_simulator(theta)
  File "/home/hananel/.local/lib/python3.8/site-packages/sbi/inference/base.py", line 133, in <lambda>
    self._batched_simulator = lambda theta: simulate_in_batches(
  File "/home/hananel/.local/lib/python3.8/site-packages/sbi/simulators/simutils.py", line 55, in simulate_in_batches
    simulation_outputs = Parallel(n_jobs=num_workers)(
  File "/home/hananel/.local/lib/python3.8/site-packages/joblib/parallel.py", line 1042, in __call__
    self.retrieve()
  File "/home/hananel/.local/lib/python3.8/site-packages/joblib/parallel.py", line 921, in retrieve
    self._output.extend(job.get(timeout=self.timeout))
  File "/home/hananel/.local/lib/python3.8/site-packages/joblib/_parallel_backends.py", line 542, in wrap_future_result
    return future.result(timeout=timeout)
  File "/usr/lib/python3.8/concurrent/futures/_base.py", line 439, in result
    return self.__get_result()
  File "/usr/lib/python3.8/concurrent/futures/_base.py", line 388, in __get_result
    raise self._exception
_pickle.PicklingError: Could not pickle the task to send it to the workers.
Running 100 simulations in 100 batches.:  19%|█▉        | 19/100 [00:00<00:01, 65.10it/s]

[michaeldeistler]

Hi,

thanks for reporting this issue! I need some more details: what version of sbi and neuron are you using and what operating system are you working on?

Also, could I ask you to try (completely outside of sbi, once you have defined your simulator):

import pickle
with open('filename.pickle', 'wb') as handle:
    pickle.dump(your_simulator, handle, protocol=pickle.HIGHEST_PROTOCOL)

as well as

import dill as pickle
with open('filename.pickle', 'wb') as handle:
    pickle.dump(your_simulator, handle, protocol=pickle.HIGHEST_PROTOCOL)

and let me know which ones of them throw an error and (possibly) which ones worked.

Michael

[Hananel-Hazan]

You are right, sorry for not providing those details.
I am using:

• Red Hat Enterprise Linux 7 and Ubuntu 20.04.1 LTS under windows subsystem with GPU support.
• sbi version is 0.12.1
• Neuron version is 7.8.1.1

Both of the pickle methods run successfully. However, I am not sure what you meant by your_simulator, Neuron is written in C that has a Python interface. So my simulator is only wrapper functions that run C program.

For your convenience, below is the sample of the code that I am using that reproduces the error.

Neuron file HHNeuron.hoc

ilow = 20   	// starting injection current (pA)
istep = 10		// current step increment (pA)
nsteps = 35		// number of steps to run


begintemplate HH

	public isoma
	create isoma
	objectvar stim

	proc init() {
		iamp = $1	
		isoma {
				insert hh			
		}
	}
endtemplate HH

objectvar HHcell[nsteps]
objref data_vecs_HH[nsteps]

proc initHH(){
	for i=0, nsteps-1 {	 
		HHcell[i] = new HH(ilow+(istep*i))  	
		data_vecs_HH[i] = new Vector()
	}
}

proc runHHsim() {	
	dt = 0.01
	v_init = -60
	tstop = 300
	tctr = 0
	
	finitialize(v_init)
		
	while(t<tstop) {
		if (tctr % 10 == 0) {
			for i=0, nsteps-1 {
				data_vecs_HH[i].append(HHcell[i].isoma.v(0.5))
			}
		}
		tctr = tctr + 1
		fadvance()
	}	
}

Python Script

import numpy as np
import torch
import pickle

# sbi
import sbi.utils as utils
from sbi.inference.base import infer

# Neuron
from neuron import h

def runHH(param):
    x, y, z = param

    h.initHH()
    # change the values of the model in Neuron
    for c in range(h.HHcell.__len__()):
        h.HHcell[c].isoma.gl_hh = x
        h.HHcell[c].isoma.gkbar_hh = y
        h.HHcell[c].isoma.gnabar_hh = z
    #  Run the model
    h.runHHsim()

    out_voltage = None
    n_shape = [h.data_vecs_HH.__len__(), h.data_vecs_HH[0].__len__()]
    if n_shape[0] > 0:
        out_voltage = torch.zeros(n_shape, dtype=torch.float)
        for y in range(n_shape[0]):
            out_voltage[y, :] = torch.as_tensor(h.data_vecs_HH[y].as_numpy())
    return out_voltage
	
def test_parameters(params):
    if params.dim() == 1:
        params = params.unsqueeze(0)
    n_tests = params.shape[0]

    Error = torch.zeros((n_tests, 315))
    for iter in range(n_tests):
        Error[iter, :] = calculate_summary_statistics(
            runHH(params[iter, :])
        )

    return Error
    
def calculate_summary_statistics(voltTrace):   
    return(torch.rand(315))
    
# Neuron Sim
h.load_file("stdrun.hoc")  # for run control
h.load_file("HHneuron.hoc")  # run the model
    
#             gl,  gK, gNa
prior_min = [0.00001, 0.001, 0.01]  # fields[0], fields[1], ....
prior_max = [0.001, 0.1, 1.0]  # fields[0], fields[1], ....
num_posterior_draw = 10
num_simulations = 30

prior = utils.torchutils.BoxUniform(low=torch.as_tensor(prior_min), high=torch.as_tensor(prior_max))
infer_obj = infer(test_parameters, prior, method='SNPE', num_simulations=num_simulations, num_workers=2)

[michaeldeistler]

That's great, thanks for the minimal code example! We will have a look and get back to you (I'm busy this week, but I'll find some time next week).

Michael

[michaeldeistler]

Hi, we have just added a solution to your problem here.

Feel free to re-open this issue if the above does not work.

Aside, using your simulator, I experienced issues because you create torch.tensors that might lie on the GPU, and this quickly filled up memory when using multiprocessing. You might want to either enforce using cpu or just resort to numpy for large tensors:

def test_parameters(params):
    import numpy as np
    import torch
    from neuron import h

    # Neuron Sim
    h.load_file("stdrun.hoc")  # for run control
    h.load_file("HHNeuron.hoc")  # run the model

    x, y, z = params

    h.initHH()
    # change the values of the model in Neuron
    for c in range(h.HHcell.__len__()):
        h.HHcell[c].isoma.gl_hh = x
        h.HHcell[c].isoma.gkbar_hh = y
        h.HHcell[c].isoma.gnabar_hh = z
    #  Run the model
    h.runHHsim()

    n_shape = [h.data_vecs_HH.__len__(), h.data_vecs_HH[0].__len__()]
    # a = torch.ones(10000, 10000)
    if n_shape[0] > 0:
        out_voltage = np.empty((0, 3001))
        for y in range(n_shape[0]):  # n_shape[0]
            dummy = h.data_vecs_HH[y].as_numpy()[
                None,
            ]
            out_voltage = np.concatenate((out_voltage, dummy), axis=0)

    # Computer summary stats
    # summary stats have to be withing the simulation code or have to be imported at the beginning of the simulator.
    Error = out_voltage[3, 100:102]

    return Error

All the best
Michael

[Hananel-Hazan]

Thanks!