25 lines
816 B
Python
25 lines
816 B
Python
import tensorflow as tf
|
|
|
|
|
|
class LabelEmbedding(tf.keras.layers.Layer):
|
|
def __init__(self, embed_dim, **kwargs):
|
|
super(LabelEmbedding, self).__init__(**kwargs)
|
|
self.embed_dim = embed_dim
|
|
|
|
def build(self, input_shape):
|
|
weight_shape = [input_shape[1], self.embed_dim]
|
|
# tf 1.15 mean(0.0), std(1.0) default value of tf.initializers.random_normal()
|
|
w_init = tf.random.normal(shape=weight_shape, mean=0.0, stddev=1.0)
|
|
self.w = tf.Variable(w_init, name='w', trainable=True)
|
|
|
|
def call(self, inputs, training=None, mask=None):
|
|
x = tf.matmul(inputs, self.w)
|
|
return x
|
|
|
|
def get_config(self):
|
|
config = super(LabelEmbedding, self).get_config()
|
|
config.update({
|
|
'embed_dim': self.embed_dim,
|
|
})
|
|
return config
|