LateralDendriticInput() Takes a lot of time.

[a2l5t8]

I run a code using LateralDendriticInput() implemented in conex and I used pyinstrument to monitor it, here is the code

net = Network(behavior = prioritize_behaviors([
    TimeResolution(dt = 1)
]))

ng = NeuronGroup(net = net, size = 28*28*2, behavior = prioritize_behaviors([
    NeuronAxon(),
    SimpleDendriteStructure(),
    SimpleDendriteComputation(),
    LIF(tau = 4, R = 2, threshold = 5, v_rest = 1, v_reset = 0)
]))

sg = SynapseGroup(net = net, src = ng, dst = ng, behavior = prioritize_behaviors([
    SynapseInit(),
    WeightInitializer(weights = torch.rand(1,1,1, 25, 25)),
    SimpleDendriticInput(),
    LateralDendriticInput(),
]))


with pyinstrument.Profiler(async_mode="disabled") as p : 
    net.initialize()
    net.simulate_iterations(1000)

and here is the monitoring results

100.0% <cell line: 25>  <ipython-input-8-8dcb09eec0e0>:1
├─ 98.5% Network.simulate_iterations  pymonntorch/NetworkCore/Network.py:320
│  ├─ 98.4% Network.simulate_iteration  pymonntorch/NetworkCore/Network.py:290
│  │  ├─ 97.4% LateralDendriticInput._wrapped_call_impl  torch/nn/modules/module.py:1732
│  │  │  ├─ 97.2% tuple._call_impl  torch/nn/modules/module.py:1740
│  │  │  │  ├─ 90.4% LateralDendriticInput.forward  conex/behaviors/synapses/dendrites.py:53
│  │  │  │  │  ├─ 89.5% LateralDendriticInput.calculate_input  conex/behaviors/synapses/dendrites.py:207
│  │  │  │  │  │  ├─ 86.4% _VariableFunctionsClass.conv3d  <built-in>
│  │  │  │  │  │  ├─ 1.2% [self]  conex/behaviors/synapses/dendrites.py
│  │  │  │  │  │  ├─ 1.1% NeuronAxon.get_spike  conex/behaviors/neurons/axon.py:70
. . . . . . . . . . . . . . . . 
. . . . . . . . . . . . . . . .

Is there any way to speed up the process?

[a2l5t8]

So There are some problems with the speed of the LateralDendriticInput() however the new implementation below solves that problem almost completely

def triple_one_idx(i, v):
    res = [1, 1, 1]
    res[i] = v
    return res


class LatheralWeight2Sparse(Behavior):
    # Put this after weight initializtion timestamp
    # and use SparseDendriticInput instead of usual LateralDendriticInput
    # for dense, flip the value of sparse into False and use SimpleDendriticInput
    def __init__(self, *args, r_sparse=True, **kwargs):
        super().__init__(*args, r_sparse=r_sparse, **kwargs)
        self.r_sparse = r_sparse

    def initialize(self, sg):
        weight = sg.weights[0, 0]  # weight will be 3D
        src_shape = sg.src_shape
        src_stride = (src_shape[1] * src_shape[2], src_shape[2], 1)
        src_numel = src_stride[0] * src_shape[0]

        w_ranges = [
            torch.arange(-(x // 2), 1 + x // 2).reshape(-1, 1) for x in weight.shape
        ]
        inp_ranges = [torch.arange(x).to(torch.float64) for x in src_shape]

        dim_idx = [torch.add(w_ranges[i], inp_ranges[i]) for i in range(3)]
        for i in range(3):
            dim_idx[i][dim_idx[i] < 0] = torch.nan
            dim_idx[i][dim_idx[i] >= src_shape[i]] = torch.nan
        dim_idx = [dim_idx[i] * src_stride[i] for i in range(3)]

        src_idx = sum(
            [
                dim_idx[i].reshape(
                    *triple_one_idx(i, weight.shape[i]),
                    *triple_one_idx(i, src_shape[i])
                )
                for i in range(3)
            ]
        )
        src_idx = src_idx.reshape(weight.numel(), src_numel)

        new_w = weight.view(-1, 1).expand(-1, src_numel)
        dst_idx = torch.arange(src_numel).view(1, -1).expand(weight.numel(), -1)

        sp_idx = torch.stack(
            [
                src_idx.reshape(
                    -1,
                ),
                dst_idx.reshape(
                    -1,
                ),
            ]
        )
        new_w = new_w.reshape(
            -1,
        )

        ok_idx = ~sp_idx[0].isnan()
        sp_idx = sp_idx[:, ok_idx]
        new_w = new_w[ok_idx]
        
        sp_w = torch.sparse_coo_tensor(sp_idx, new_w, (src_numel, src_numel))
        if self.r_sparse:
            sg.weights = sp_w.to_sparse_csc()
        else:
            sg.weights = sp_w.to_dense()

        sg.weights = sg.weights.to(self.device).to(sg.def_dtype)

We tested this on a dummy network with only one NeuronGroup and one SynapseGroup having LateralDendriticInput() for 1000 iterations and it nearly took 15s for the whole simulation, however we removed LateralDendriticInput() and used this new behavior with both SimpleDendriticInput() and SparseDendriticInput() and it was done in about 0.7s !!

Huge thanks to @saeedark for the new implementation :)