🚧 rewrite classification models under tf.keras schema

BrikerMan · Jun 24, 2019 · 54e1c91 · 54e1c91
1 parent ad846ac
commit 54e1c91
Showing 1 changed file with 84 additions and 1 deletion.
diff --git a/kashgari/tasks/classification/models.py b/kashgari/tasks/classification/models.py
@@ -14,7 +14,7 @@
 from kashgari.tasks.classification.base_model import BaseClassificationModel
 
 
-class BLSTMModel(BaseClassificationModel):
+class BiLSTM_Model(BaseClassificationModel):
 
     @classmethod
     def get_default_hyper_parameters(cls) -> Dict[str, Dict[str, Any]]:
@@ -42,6 +42,89 @@ def build_model_arc(self):
         self.tf_model = tf.keras.Model(embed_model.inputs, output_tensor)
 
 
+class CNN_Model(BaseClassificationModel):
+
+    @classmethod
+    def get_default_hyper_parameters(cls) -> Dict[str, Dict[str, Any]]:
+        return {
+            'conv1d_layer': {
+                'filters': 128,
+                'kernel_size': 5,
+                'activation': 'relu'
+            },
+            'max_pool_layer': {},
+            'dense_layer': {
+                'units': 64,
+                'activation': 'relu'
+            },
+            'activation_layer': {
+                'activation': 'softmax'
+            },
+        }
+
+    def build_model_arc(self):
+        output_dim = len(self.pre_processor.label2idx)
+        config = self.hyper_parameters
+        embed_model = self.embedding.embed_model
+
+        # build model structure in sequent way
+        layers_seq = []
+        layers_seq.append(L.Conv1D(**config['conv1d_layer']))
+        layers_seq.append(L.GlobalMaxPooling1D(**config['max_pool_layer']))
+        layers_seq.append(L.Dense(**config['dense_layer']))
+        layers_seq.append(L.Dense(output_dim, **config['activation_layer']))
+
+        tensor = embed_model.output
+        for layer in layers_seq:
+            tensor = layer(tensor)
+
+        self.tf_model = tf.keras.Model(embed_model.inputs, tensor)
+
+
+class CNN_LSTM_Model(BaseClassificationModel):
+
+    @classmethod
+    def get_default_hyper_parameters(cls) -> Dict[str, Dict[str, Any]]:
+        return {
+            'conv_layer': {
+                'filters': 32,
+                'kernel_size': 3,
+                'padding': 'same',
+                'activation': 'relu'
+            },
+            'max_pool_layer': {
+                'pool_size': 2
+            },
+            'lstm_layer': {
+                'units': 100
+            },
+            'activation_layer': {
+                'activation': 'softmax'
+            },
+        }
+
+    def build_model_arc(self):
+        output_dim = len(self.pre_processor.label2idx)
+        config = self.hyper_parameters
+        embed_model = self.embedding.embed_model
+
+        layers_seq = []
+        layers_seq.append(L.Conv1D(**config['conv_layer']))
+        layers_seq.append(L.MaxPooling1D(**config['max_pool_layer']))
+        layers_seq.append(L.LSTM(**config['lstm_layer']))
+        layers_seq.append(L.Dense(output_dim, **config['activation_layer']))
+
+        tensor = embed_model.output
+        for layer in layers_seq:
+            tensor = layer(tensor)
+
+        self.tf_model = tf.keras.Model(embed_model.inputs, tensor)
+
+
+BLSTMModel = BiLSTM_Model
+CNNModel = CNN_Model
+CNNLSTMModel = CNN_LSTM_Model
+
 if __name__ == "__main__":
     print(BLSTMModel.get_default_hyper_parameters())
     logging.basicConfig(level=logging.DEBUG)