ma_cisco_malware/models/__init__.py

160 lines
5.7 KiB
Python
Raw Normal View History

import keras.backend as K
from keras.models import Model
from models import deep1
from models.renes_networks import selu
from . import flat_2, pauls_networks, renes_networks
def create_model(model, output_type):
if output_type == "both":
return Model(inputs=[model.in_domains, model.in_flows], outputs=(model.out_client, model.out_server))
elif output_type == "client":
return Model(inputs=[model.in_domains, model.in_flows], outputs=(model.out_client,))
else:
raise Exception("unknown model output")
def get_models_by_params(params: dict):
# decomposing param section
# mainly embedding model
network_type = params.get("type")
network_depth = params.get("depth")
embedding_size = params.get("embedding")
filter_embedding = params.get("filter_embedding")
kernel_embedding = params.get("kernel_embedding")
hidden_embedding = params.get("dense_embedding")
2017-10-05 15:26:53 +02:00
# dropout = params.get("dropout")
# mainly prediction model
flow_features = params.get("flow_features")
window_size = params.get("window_size")
domain_length = params.get("domain_length")
filter_main = params.get("filter_main")
kernel_main = params.get("kernel_main")
dense_dim = params.get("dense_main")
model_output = params.get("model_output", "both")
# create models
if network_depth == "flat1":
networks = pauls_networks
elif network_depth == "flat2":
networks = flat_2
elif network_depth == "deep1":
2017-10-09 15:10:15 +02:00
networks = deep1
elif network_depth == "deep2":
networks = renes_networks
else:
raise ValueError("network not found")
2017-11-05 22:52:50 +01:00
domain_cnn = networks.get_embedding(embedding_size, domain_length, filter_embedding, kernel_embedding,
hidden_embedding, 0.5)
if network_type == "final":
model = networks.get_model(0.25, flow_features, window_size, domain_length,
filter_main, kernel_main, dense_dim, domain_cnn)
model = create_model(model, model_output)
elif network_type in ("inter", "staggered"):
model = networks.get_new_model(0.25, flow_features, window_size, domain_length,
filter_main, kernel_main, dense_dim, domain_cnn)
model = create_model(model, model_output)
elif network_type == "long":
model = networks.get_new_model2(0.25, flow_features, window_size, domain_length,
filter_main, kernel_main, dense_dim, domain_cnn)
model = create_model(model, model_output)
elif network_type == "soft":
model = networks.get_new_soft(0.25, flow_features, window_size, domain_length,
filter_main, kernel_main, dense_dim, domain_cnn)
model = create_model(model, model_output)
conv_server = model.get_layer("conv_server").trainable_weights
conv_client = model.get_layer("conv_client").trainable_weights
l1 = [0.001 * K.sum(K.abs(x - y)) for (x, y) in zip(conv_server, conv_client)]
model.add_loss(l1)
dense_server = model.get_layer("dense_server").trainable_weights
dense_client = model.get_layer("dense_client").trainable_weights
l2 = [0.001 * K.sum(K.abs(x - y)) for (x, y) in zip(dense_server, dense_client)]
model.add_loss(l2)
else:
raise ValueError("network type not found")
2017-11-06 21:51:49 +01:00
return model
2017-10-05 15:26:53 +02:00
def get_server_model_by_params(params: dict):
# decomposing param section
# mainly embedding model
network_depth = params.get("depth")
embedding_size = params.get("embedding")
input_length = params.get("input_length")
filter_embedding = params.get("filter_embedding")
kernel_embedding = params.get("kernel_embedding")
hidden_embedding = params.get("dense_embedding")
# mainly prediction model
flow_features = params.get("flow_features")
domain_length = params.get("domain_length")
dense_dim = params.get("dense_main")
# create models
if network_depth == "flat1":
networks = pauls_networks
elif network_depth == "flat2":
networks = flat_2
elif network_depth == "deep1":
networks = deep1
elif network_depth == "deep2":
2017-10-05 15:26:53 +02:00
networks = renes_networks
else:
raise Exception("network not found")
embedding_model = networks.get_embedding(embedding_size, input_length, filter_embedding, kernel_embedding,
hidden_embedding, 0.5)
return networks.get_server_model(flow_features, domain_length, dense_dim, embedding_model)
def get_custom_objects():
return dict([
("precision", precision),
("recall", recall),
("f1_score", f1_score),
("selu", selu)
])
def get_metric_functions():
return [precision, recall, f1_score]
def precision(y_true, y_pred):
# Count positive samples.
true_positives = K.sum(K.round(K.clip(y_true * y_pred, 0, 1)))
predicted_positives = K.sum(K.round(K.clip(y_pred, 0, 1)))
return true_positives / (predicted_positives + K.epsilon())
def recall(y_true, y_pred):
# Count positive samples.
true_positives = K.sum(K.round(K.clip(y_true * y_pred, 0, 1)))
possible_positives = K.sum(K.round(K.clip(y_true, 0, 1)))
return true_positives / (possible_positives + K.epsilon())
def f1_score(y_true, y_pred):
return f_score(1)(y_true, y_pred)
def f05_score(y_true, y_pred):
return f_score(0.5)(y_true, y_pred)
def f_score(beta):
def _f(y_true, y_pred):
p = precision(y_true, y_pred)
r = recall(y_true, y_pred)
bb = beta ** 2
fbeta_score = (1 + bb) * (p * r) / (bb * p + r + K.epsilon())
return fbeta_score
return _f