Керас: Есть ли обходной путь для разделения вывода промежуточного слоя без использования слоя Lamda?

Скажем, у меня есть 10x10x4 промежуточный выход сверточного слоя, который мне нужно разделить на 100 1x1x4 томов и применить softmax к каждому, чтобы получить 100 выходов из сети. Есть ли способ сделать это без использования слоя Lambda? Проблема со слоем Lambda в этом случае заключается в том, что эта простая задача разделения занимает 100 проходов через слой лямбда во время прямого прохода, что сильно снижает производительность сети для моего практического использования. Предложите более быстрый способ сделать это.

Изменить: я уже пробовал подход Softmax + Reshape, прежде чем задавать вопрос. При таком подходе я бы получил матрицу 10x10x4, преобразованную в тензор 100x4 с использованием Reshape в качестве вывода. Что мне действительно нужно, так это сеть с несколькими выходами и 100 различными выходами. В моем приложении невозможно совместно оптимизировать матрицу 10x10, но я получаю хорошие результаты, используя сеть со 100 различными выходами со слоем Lambda.

Вот фрагменты кода моего подхода с использованием функционального API Keras:

Со слоем Lambda (медленно, дает 100 тензоров формы (None, 4) по желанию):

# Assume conv_output is output from a convolutional layer with shape (None, 10, 10,4)
preds = []
for i in range(10):
    for j in range(10):
        y = Lambda(lambda x, i,j: x[:, i, j,:],  arguments={'i': i,'j':j})(conv_output)
        preds.append(Activation('softmax',name='predictions_' + str(i*10+j))(y))

model = Model(inputs=img, outputs=preds, name='model')
model.compile(loss='categorical_crossentropy',
              optimizer=Adam(),
              metrics=['accuracy']

С Softmax + Reshape (быстро, но дает тензор формы (None, 100, 4))

# Assume conv_output is output from a convolutional layer with shape (None, 10, 10,4)
y = Softmax(name='softmax', axis=-1)(conv_output)
preds = Reshape([100, 4])(y)
model = Model(inputs=img, outputs=preds, name='model')
model.compile(loss='categorical_crossentropy',
              optimizer=Adam(),
              metrics=['accuracy']

Я не думаю, что во втором случае можно индивидуально оптимизировать каждый из 100 выходов (вероятно, можно думать об этом как об изучении совместного распределения, тогда как мне нужно изучить маргиналы, как в первом случае). Пожалуйста, дайте мне знать, есть ли способ быстрее выполнить то, что я делаю со слоем Lambda в первом фрагменте кода.