from collections import namedtuple import keras import keras.backend as K from keras.engine import Input, Model as KerasModel from keras.engine.topology import Layer from keras.layers import Conv1D, Dense, Dropout, Embedding, GlobalAveragePooling1D, GlobalMaxPooling1D, TimeDistributed import dataset Model = namedtuple("Model", ["in_domains", "in_flows", "out_client", "out_server"]) def get_domain_embedding_model(embedding_size, input_length, filter_size, kernel_size, hidden_dims, drop_out=0.5) -> KerasModel: x = y = Input(shape=(input_length,)) y = Embedding(input_dim=dataset.get_vocab_size(), output_dim=embedding_size)(y) y = Conv1D(filter_size, kernel_size, activation='relu')(y) y = GlobalMaxPooling1D()(y) y = Dropout(drop_out)(y) y = Dense(hidden_dims, activation="relu")(y) return KerasModel(x, y) def get_domain_embedding_model2(embedding_size, input_length, filter_size, kernel_size, hidden_dims, drop_out=0.5) -> KerasModel: x = y = Input(shape=(input_length,)) y = Embedding(input_dim=dataset.get_vocab_size(), output_dim=embedding_size)(y) y = Conv1D(filter_size, kernel_size, activation='relu')(y) y = Conv1D(filter_size, kernel_size, activation='relu')(y) y = Conv1D(filter_size, kernel_size, activation='relu')(y) y = GlobalAveragePooling1D()(y) y = Dense(hidden_dims, activation="relu")(y) return KerasModel(x, y) def get_final_model(cnnDropout, flow_features, window_size, domain_length, cnn_dims, kernel_size, dense_dim, cnn) -> Model: ipt_domains = Input(shape=(window_size, domain_length), name="ipt_domains") encoded = TimeDistributed(cnn, name="domain_cnn")(ipt_domains) ipt_flows = Input(shape=(window_size, flow_features), name="ipt_flows") merged = keras.layers.concatenate([encoded, ipt_flows], -1) # CNN processing a small slides of flow windows y = Conv1D(cnn_dims, kernel_size, activation='relu')(merged) # remove temporal dimension by global max pooling y = GlobalMaxPooling1D()(y) y = Dropout(cnnDropout)(y) y = Dense(dense_dim, activation='relu')(y) out_client = Dense(1, activation='sigmoid', name="client")(y) out_server = Dense(1, activation='sigmoid', name="server")(y) return Model(ipt_domains, ipt_flows, out_client, out_server) def get_inter_model(dropout, flow_features, window_size, domain_length, cnn_dims, kernel_size, dense_dim, cnn) -> Model: ipt_domains = Input(shape=(window_size, domain_length), name="ipt_domains") ipt_flows = Input(shape=(window_size, flow_features), name="ipt_flows") encoded = TimeDistributed(cnn, name="domain_cnn")(ipt_domains) merged = keras.layers.concatenate([encoded, ipt_flows], -1) y = Dense(dense_dim, activation="relu", name="dense_server")(merged) out_server = Dense(1, activation="sigmoid", name="server")(y) merged = keras.layers.concatenate([merged, y], -1) # CNN processing a small slides of flow windows y = Conv1D(cnn_dims, kernel_size, activation='relu')(merged) # remove temporal dimension by global max pooling y = GlobalMaxPooling1D()(y) y = Dropout(dropout)(y) y = Dense(dense_dim, activation='relu', name="dense_client")(y) out_client = Dense(1, activation='sigmoid', name="client")(y) return Model(ipt_domains, ipt_flows, out_client, out_server) def get_server_model(flow_features, domain_length, dense_dim, cnn): ipt_domains = Input(shape=(domain_length,), name="ipt_domains") ipt_flows = Input(shape=(flow_features,), name="ipt_flows") encoded = cnn(ipt_domains) cnn.name = "domain_cnn" merged = keras.layers.concatenate([encoded, ipt_flows], -1) y = Dense(dense_dim, activation="relu", name="dense_server")(merged) out_server = Dense(1, activation="sigmoid", name="server")(y) return KerasModel(inputs=[ipt_domains, ipt_flows], outputs=out_server) def get_long_model(dropout, flow_features, window_size, domain_length, cnn_dims, kernel_size, dense_dim, cnn) -> Model: ipt_domains = Input(shape=(window_size, domain_length), name="ipt_domains") ipt_flows = Input(shape=(window_size, flow_features), name="ipt_flows") encoded = TimeDistributed(cnn, name="domain_cnn")(ipt_domains) merged = keras.layers.concatenate([encoded, ipt_flows], -1) y = Conv1D(cnn_dims, kernel_size, activation='relu', name="conv_server")(merged) # remove temporal dimension by global max pooling y = GlobalMaxPooling1D()(y) y = Dropout(dropout)(y) y = Dense(dense_dim, activation="relu", name="dense_server")(y) out_server = Dense(1, activation="sigmoid", name="server")(y) # CNN processing a small slides of flow windows y = Conv1D(cnn_dims, kernel_size, activation='relu', name="conv_client")(merged) # remove temporal dimension by global max pooling y = GlobalMaxPooling1D()(y) y = Dropout(dropout)(y) y = Dense(dense_dim, activation='relu', name="dense_client")(y) out_client = Dense(1, activation='sigmoid', name="client")(y) return Model(ipt_domains, ipt_flows, out_client, out_server) class CrossStitch2(Layer): def __init__(self, **kwargs): super(CrossStitch2, self).__init__(**kwargs) def build(self, input_shape): # Create a trainable weight variable for this layer. self.s = self.add_weight(name='cross-stitch-s', shape=(1,), initializer='uniform', trainable=True) self.d = self.add_weight(name='cross-stitch-d', shape=(1,), initializer='uniform', trainable=True) super(CrossStitch2, self).build(input_shape) def call(self, xs): x1, x2 = xs out = x1 * self.s + x2 * self.d print("==>", x1, x2, out) return out def compute_output_shape(self, input_shape): return input_shape[0] class CrossStitchMix2(Layer): def __init__(self, **kwargs): super(CrossStitchMix2, self).__init__(**kwargs) def build(self, input_shape): # Create a trainable weight variable for this layer. self.s = self.add_weight(name='cross-stitch-s', shape=(1,), initializer='uniform', trainable=True) self.d = self.add_weight(name='cross-stitch-d', shape=(1,), initializer='uniform', trainable=True) super(CrossStitchMix2, self).build(input_shape) def call(self, xs): x1, x2 = xs out = (x1 * self.s, x2 * self.d) out = K.concatenate(out, axis=-1) return out def compute_output_shape(self, input_shape): return (input_shape[0][0], input_shape[0][1] + input_shape[1][1]) def get_sluice_model(dropout, flow_features, window_size, domain_length, cnn_dims, kernel_size, dense_dim, cnn) -> Model: ipt_domains = Input(shape=(window_size, domain_length), name="ipt_domains") ipt_flows = Input(shape=(window_size, flow_features), name="ipt_flows") encoded = TimeDistributed(cnn, name="domain_cnn")(ipt_domains) merged = keras.layers.concatenate([encoded, ipt_flows], -1) y1 = Conv1D(cnn_dims, kernel_size, activation='relu', name="conv_server")(merged) y1 = GlobalMaxPooling1D()(y1) y2 = Conv1D(cnn_dims, kernel_size, activation='relu', name="conv_client")(merged) y2 = GlobalMaxPooling1D()(y2) c11 = CrossStitch2()([y1, y2]) c12 = CrossStitch2()([y1, y2]) y1 = Dropout(dropout)(c11) y1 = Dense(dense_dim, activation="relu", name="dense_server")(y1) y2 = Dropout(dropout)(c12) y2 = Dense(dense_dim, activation='relu', name="dense_client")(y2) c21 = CrossStitch2()([y1, y2]) c22 = CrossStitch2()([y1, y2]) beta1 = CrossStitchMix2()([c11, c21]) beta2 = CrossStitchMix2()([c12, c22]) out_server = Dense(1, activation="sigmoid", name="server")(beta1) out_client = Dense(1, activation='sigmoid', name="client")(beta2) return Model(ipt_domains, ipt_flows, out_client, out_server)