keras实现方式中的稀疏连接层太慢

Question

我试图建立一个稀疏连接的层，其中每个输出神经元连接到正好n个输入

我已经想出了一个似乎可行的实现。当我用自定义层替换稠密层时，我的模型似乎会收敛。我通常会添加更多的神经元来补偿由于限制连接数量而导致的预测能力的损失

但是，实现速度非常慢（比标准密集层慢30倍）。你知道我在哪里可以改进我的代码使其“工作”吗？还是另一种会导致连接层稀疏的方法

 class NConnected(tf.keras.layers.Layer):

    #creator
    def __init__(self, units=32, n_max = 5):
        super(NConnected, self).__init__()
        self.units = units
        self.n_max = n_max
        
   #creates weights
    def build(self, input_shape):
        self.w = self.add_weight(name='w', 
                                  shape=(self.n_max,self.units),
                                  initializer='uniform',
                                  trainable=True)

        self.b = self.add_weight(name='b', 
                                   shape=self.units,
                                   initializer='zeros',
                                   trainable=True)
        
        mask_t = []

        for i in range(self.units):
            
            
            a = range(1,input_shape[-1]+1)
            r = random.sample(a,self.n_max)

            mask_t.append(np.array([i in r for i in a]))

        mask_t = tf.constant(np.array(mask_t))
        
        self.mask = mask_t

        self.built=True

   #operation:
    def call(self, inputs):
        
        m = tf.map_fn(fn=lambda t: tf.boolean_mask(inputs,t,axis=1),elems = self.mask, fn_output_signature=tf.float32)
        m = tf.transpose(m, [1, 0, 2])
        res = tf.math.reduce_sum(tf.tensordot(m,self.w,axes=1),axis=2)+self.b
        
        return res

   #for saving the model - only necessary if you have parameters in __init__
    def get_config(self):
        config = super(NConnected, self).get_config()
        return config

如前所述，该层比同等大小的致密层慢得多。以简单的向前传球为例：

小批量随机数据：

x_test = tf.random.normal((1000,100), mean=0.0, stddev=1.0, dtype=tf.dtypes.float32, seed=None, name=None)

简单测试循环：

%%time

test_n = NConnected(units=32, n_max = 3)

for i in range(100):
    test_n(x_test)

返回：

CPU时间：用户4.77秒，系统59.1毫秒，总计4.83秒 壁时间：4.72秒

而：

%%time
    
test_dense = layers.Dense(units=32)
    
for i in range(100):
   test_dense(x_test)

返回：

CPU时间：用户47.9毫秒，系统7.2毫秒，总计55.1毫秒 墙壁时间：45.3毫秒

所以大约慢100倍，这使得在实践中几乎不可能使用稀疏层