Создание пользовательского слоя и его свертка

Question

Рассмотрим код:

from tensorflow.contrib import layers

#input_features is given, and is 1x20x20x1000 for instance.
N_outputs = 200
N_intermediate = 50

out_lst = []
for i in xrange(N_outputs):
    net = layers.conv2d(
        input_features,
        N_intermediate, [1, 1]
    )
    out = layers.conv2d(
        net,
        1, [1, 1],
        activation_fn=None,
        normalizer_fn=None
    )

    out_lst.append(out)

logits = tf.concat(out_lst,axis = 3)

В целом, сеть должна принимать входные данные размером 1x20x20x1000 и возвращать выходные данные размером 1x20x20x200, где каждый из 200 выходных логитов проходит через свой собственный единственный размерный слойN_intermediate.

Проблема заключается в том, что построение графика происходит довольно медленно, а затем требует слишком много памяти, что приводит к ошибке памяти.Для сравнения, следующий код

#input_features is given, and is 1x20x20x1000. 
N_outputs = 200
N_intermediate = 50

net = layers.conv2d(
        input_features,
        N_intermediate*N_outputs, [1, 1]
    )
net = layers.conv2d(
        net,
        N_outputs, [1, 1],
        activation_fn=None,
        normalizer_fn=None
)

logits = net

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

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