rene/ma_cisco_malware
1
0
Fork 0
Code Issues Pull Requests Releases Wiki Activity

add argument for model outputs

Browse Source 
BUG: need to check --new_model --model_output server
This commit is contained in:
René Knaebel 2017-08-02 12:58:09 +02:00
parent f4da147688
commit 1e781d5491
5 changed files with 45 additions and 15 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
arguments.py
View File

@ -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)

12
main.py
View File

@ -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,

7
models/__init__.py
View File

@ -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

19
models/pauls_networks.py
View File

@ -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,))

19
models/renes_networks.py
View File

@ -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,))