| 5 | |
| 6 | |
| 7 | class MyDense(tf.keras.layers.Layer): |
| 8 | def __init__(self, units, activation=None, **kwargs): |
| 9 | super(MyDense, self).__init__(**kwargs) |
| 10 | self.units = units |
| 11 | self.w = None |
| 12 | self.activation = activation |
| 13 | |
| 14 | def build(self, input_shape): |
| 15 | self.w = self.add_weight( |
| 16 | name="weight", shape=(input_shape[1], self.units), trainable=True |
| 17 | ) |
| 18 | super(MyDense, self).build(input_shape) |
| 19 | |
| 20 | def call(self, inputs): |
| 21 | if self.activation == None: |
| 22 | return sgx.e_matmul(inputs, self.w) |
| 23 | elif self.activation == "relu": |
| 24 | return sgx.e_relu(sgx.e_matmul(inputs, self.w)) |
| 25 | elif self.activation == "sigmoid": |
| 26 | return sgx.e_sigmoid(sgx.e_matmul(inputs, self.w)) |
| 27 | |
| 28 | def compute_output_shape(self, input_shape): |
| 29 | return (input_shape[0], self.units) |
| 30 | |
| 31 | def get_config(self): |
| 32 | config = super(MyDense, self).get_config() |
| 33 | config.update({"units": self.units}) |
| 34 | config.update({"activation": self.activation}) |
| 35 | return config |
| 36 | |
| 37 | |
| 38 | class Embed_layer(Layer): |