Как преобразовать этот код из Keras в Tensorflow?

Question

Я пытаюсь преобразовать код из Keras в тензор потока, я не очень разбираюсь в Keras api, я пользователь Tensorflow, вот код Keras:

rawmeta = layers.Input(shape=(1,), dtype="string")
emb = elmolayer()(rawmeta)
d1 = layers.Dense(256, activation='relu')(emb)
yhat = layers.Dense(31, activation='softmax', name = "output_node")(d1)
model = Model(inputs=[rawmeta], outputs=yhat)
model.compile(loss='categorical_crossentropy', optimizer='adam', metrics=['accuracy'])

Где elmolayer определен следующим образом:

class elmolayer(Layer):

    def __init__(self, **kwargs):
        self.dimensions = 1024
        self.trainable=True
        super(elmolayer, self).__init__(**kwargs)

    def build(self, input_shape):
        self.elmo = hub.Module('https://tfhub.dev/google/elmo/2', trainable=self.trainable,
                               name="{}_module".format(self.name))

        self.trainable_weights += K.tf.trainable_variables(scope="^{}_module/.*".format(self.name))
        super(elmolayer, self).build(input_shape)

    def call(self, x, mask=None):
        result = self.elmo(K.squeeze(K.cast(x, tf.string), axis=1),
                      as_dict=True,
                      signature='default',
                      )['default']
        return result

    def compute_mask(self, inputs, mask=None):
        return K.not_equal(inputs, '--PAD--')

    def compute_output_shape(self, input_shape):
        return (input_shape[0], self.dimensions)

Моя реализация Tensorflow этого кода:

class Base_model(object):

    def __init__(self, elmo_embedding_matrix):

        tf.reset_default_graph()

        # define placeholders
        sentences        = tf.placeholder(tf.int32, [None, None], name='sentences')
        y_true     = tf.placeholder(tf.int32, [None, None], name='labels'  )





        self.elmo     = tf.get_variable(name="relation_embedding", shape=[elmo_embedding_matrix.shape[0],elmo_embedding_matrix.shape[1]],
                                         initializer=tf.constant_initializer(np.array(elmo_embedding_matrix)), 
                                         trainable=True,dtype=tf.float32)

        embedding_lookup   = tf.nn.embedding_lookup(self.elmo,sentences)

        d1 = tf.layers.dense(embedding_lookup, 256, tf.nn.relu)

        y_pred = tf.layers.dense(d1, 31, tf.nn.softmax)


        matches = tf.equal(tf.argmax(y_pred,1),tf.argmax(y_true,1))
        acc = tf.reduce_mean(tf.cast(matches,tf.float32))

        cross_entropy = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits_v2(labels=y_true,logits=y_pred))
        train = tf.train.AdamOptimizer(learning_rate=0.001).minimize(cross_entropy)

Моя путаница - последний плотный слой в модели keras:

yhat = layers.Dense(31, activation='softmax', name = "output_node")(d1)

В то время как вкод тензорного потока, если я использую tf.nn.softmax_cross_entropy_with_logits_v2, тогда я должен передать второй плотный слой в softmax, например,

y_pred = tf.layers.dense(d1, 31, tf.nn.softmax)

Поскольку, если я здесь использую softmax, тогда tf.nn.softmax_cross_entropy_with_logits_v2 будет снова использовать softmax для логитов.

Как преобразовать этот код Keras в Tensorflow?