diff --git a/arguments.py b/arguments.py index 9fab93f..094c602 100644 --- a/arguments.py +++ b/arguments.py @@ -22,6 +22,9 @@ parser.add_argument("--model", action="store", dest="model_path", parser.add_argument("--type", action="store", dest="model_type", default="paul") +parser.add_argument("--model_output", action="store", dest="model_output", + default="both") + parser.add_argument("--batch", action="store", dest="batch_size", default=64, type=int) diff --git a/main.py b/main.py index 8ec4d45..f723346 100644 --- a/main.py +++ b/main.py @@ -72,7 +72,8 @@ PARAMS = { 'hidden_embedding': args.domain_embedding, 'kernel_embedding': 3, 'kernels_main': 3, - 'input_length': 40 + 'input_length': 40, + 'model_output': args.model_output } @@ -159,8 +160,15 @@ def main_train(param=None): loss='binary_crossentropy', metrics=['accuracy'] + custom_metrics) + if args.model_output == "both": + labels = [client_tr, server_tr] + elif args.model_output == "client": + labels = [client_tr] + elif args.model_output == "server": + labels = [server_tr] + model.fit([domain_tr, flow_tr], - [client_tr, server_tr], + labels, batch_size=args.batch_size, epochs=args.epochs, callbacks=callbacks, diff --git a/models/__init__.py b/models/__init__.py index c6d883a..d71715c 100644 --- a/models/__init__.py +++ b/models/__init__.py @@ -22,16 +22,17 @@ def get_models_by_params(params: dict): filter_main = params.get("filter_main") kernel_main = params.get("kernels_main") dense_dim = params.get("dense_main") + model_output = params.get("model_output", "both") # create models networks = renes_networks if network_type == "rene" else pauls_networks embedding_model = networks.get_embedding(embedding_size, input_length, filter_embedding, kernel_embedding, - hidden_embedding, drop_out=dropout) + hidden_embedding, dropout) predict_model = networks.get_model(dropout, flow_features, domain_features, window_size, domain_length, - filter_main, kernel_main, dense_dim, embedding_model) + filter_main, kernel_main, dense_dim, embedding_model, model_output) new_model = networks.get_new_model(dropout, flow_features, domain_features, window_size, domain_length, - filter_main, kernel_main, dense_dim, embedding_model) + filter_main, kernel_main, dense_dim, embedding_model, model_output) return embedding_model, predict_model, new_model diff --git a/models/pauls_networks.py b/models/pauls_networks.py index 45f6aa5..bb824fd 100644 --- a/models/pauls_networks.py +++ b/models/pauls_networks.py @@ -38,7 +38,7 @@ def get_embedding(embedding_size, input_length, filter_size, kernel_size, hidden def get_model(cnnDropout, flow_features, domain_features, window_size, domain_length, cnn_dims, kernel_size, - dense_dim, cnn): + dense_dim, cnn, model_output="both"): ipt_domains = Input(shape=(window_size, domain_length), name="ipt_domains") encoded = TimeDistributed(cnn)(ipt_domains) ipt_flows = Input(shape=(window_size, flow_features), name="ipt_flows") @@ -55,11 +55,16 @@ def get_model(cnnDropout, flow_features, domain_features, window_size, domain_le y1 = Dense(1, activation='sigmoid', name="client")(y) y2 = Dense(1, activation='sigmoid', name="server")(y) - return Model(inputs=[ipt_domains, ipt_flows], outputs=(y1, y2)) + if model_output == "both": + return Model(inputs=[ipt_domains, ipt_flows], outputs=(y1, y2)) + elif model_output == "client": + return Model(inputs=[ipt_domains, ipt_flows], outputs=(y1,)) + elif model_output == "server": + return Model(inputs=[ipt_domains, ipt_flows], outputs=(y2,)) def get_new_model(dropout, flow_features, domain_features, window_size, domain_length, cnn_dims, kernel_size, - dense_dim, cnn): + dense_dim, cnn, model_output="both"): 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)(ipt_domains) @@ -77,6 +82,10 @@ def get_new_model(dropout, flow_features, domain_features, window_size, domain_l y = Dense(dense_dim, activation='relu')(y) y1 = Dense(1, activation='sigmoid', name="client")(y) - model = Model(inputs=[ipt_domains, ipt_flows], outputs=(y1, y2)) - return model + if model_output == "both": + return Model(inputs=[ipt_domains, ipt_flows], outputs=(y1, y2)) + elif model_output == "client": + return Model(inputs=[ipt_domains, ipt_flows], outputs=(y1,)) + elif model_output == "server": + return Model(inputs=[ipt_domains, ipt_flows], outputs=(y2,)) diff --git a/models/renes_networks.py b/models/renes_networks.py index a9c3af9..1f485c1 100644 --- a/models/renes_networks.py +++ b/models/renes_networks.py @@ -18,7 +18,7 @@ def get_embedding(embedding_size, input_length, filter_size, kernel_size, hidden def get_model(cnnDropout, flow_features, domain_features, window_size, domain_length, cnn_dims, kernel_size, - dense_dim, cnn): + dense_dim, cnn, model_output="both"): ipt_domains = Input(shape=(window_size, domain_length), name="ipt_domains") encoded = TimeDistributed(cnn)(ipt_domains) ipt_flows = Input(shape=(window_size, flow_features), name="ipt_flows") @@ -38,11 +38,16 @@ def get_model(cnnDropout, flow_features, domain_features, window_size, domain_le y1 = Dense(1, activation='sigmoid', name="client")(y) y2 = Dense(1, activation='sigmoid', name="server")(y) - return Model(inputs=[ipt_domains, ipt_flows], outputs=(y1, y2)) + if model_output == "both": + return Model(inputs=[ipt_domains, ipt_flows], outputs=(y1, y2)) + elif model_output == "client": + return Model(inputs=[ipt_domains, ipt_flows], outputs=(y1,)) + elif model_output == "server": + return Model(inputs=[ipt_domains, ipt_flows], outputs=(y2,)) def get_new_model(dropout, flow_features, domain_features, window_size, domain_length, cnn_dims, kernel_size, - dense_dim, cnn): + dense_dim, cnn, model_output="both"): 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)(ipt_domains) @@ -60,6 +65,10 @@ def get_new_model(dropout, flow_features, domain_features, window_size, domain_l y = Dense(dense_dim, activation='relu')(y) y1 = Dense(1, activation='sigmoid', name="client")(y) - model = Model(inputs=[ipt_domains, ipt_flows], outputs=(y1, y2)) - return model + if model_output == "both": + return Model(inputs=[ipt_domains, ipt_flows], outputs=(y1, y2)) + elif model_output == "client": + return Model(inputs=[ipt_domains, ipt_flows], outputs=(y1,)) + elif model_output == "server": + return Model(inputs=[ipt_domains, ipt_flows], outputs=(y2,))