Исправлено количество выпадений на определенной оси

Question

Допустим, у меня есть тензор, имеющий форму [35, 4, 5, 3], который представляет собой batch_size, 4 столбца, 5 выборок из разных времен, 3 канала.

То, что я хочу сделать, - это выбрать случайным образом ровно 3 из 5 выборок из разных времен для каждого столбца, каждого канала и каждой выборки в пакете. Для выбранных 3 увеличьте их на 5/3. Сброшенные 2 установлены на 0. Это код, который я сейчас использую:

def hard_dropout(x, num_drop, axis=-1, training=False):
    if not training:
        return x
    if axis < 0:
        axis = len(x.shape) + axis
    assert axis >= 0

    shape = [int(i) for i in x.shape]
    # length of dropout mask
    length = shape[axis]
    del shape[axis]
    # number of dropout mask (for each example, column and channel)
    num_repeat = np.prod(shape)
    masks = []
    mask = np.concatenate([np.ones(length - num_drop, dtype=np.float32) * length / (length - num_drop),
                           np.zeros(num_drop, dtype=np.float32)])
    with tf.name_scope('hard_dropout'):
        for i in range(num_repeat):
            masks.append(tf.random_shuffle(mask))
        masks = tf.stack(masks, axis=0)
        # reshape and transpose the masks to the shape of input tensor
        masks = tf.reshape(masks, [*shape, length])
        permutation = list(range(0,len(shape)))
        permutation.insert(axis, len(shape))
        masks = tf.transpose(masks, permutation)
        x = masks * x
    return x

Кажется, что текущий код работает нормально, но он не очень эффективен и замедляет обучение, поэтому мне интересно, есть ли лучший способ сделать это?