class Trainer(object):
def __init__(self):
from flow import Flows
self.flow = Flows()
# Train supervised with samples and targets
def train(self, algorithm, data, feature, good_labels, config):
key = self.check_keys(data, "type", "DefaultTrainer")
trainer = Trainer.get_trainer(key, DefaultTrainer())
print "Loaded training algorithm: " + str(key) + "."
fl = trainer.train(algorithm, data, feature, good_labels, config)
if not fl:
return False
else:
self.flow.addFlows(fl)
return True
def check_keys(self, dic, key, val):
if not key in dic:
dic[key] = val
return dic[key]
@staticmethod
def get_trainer(name, default=None):
import sys
try:
return getattr(sys.modules[__name__], name)()
except Exception as e:
print "Trainer \"" + name + "\"does not exist."
return default
def trainAll(self, feature, algorithm, good_labels, manager, is_binary):
samples = self.flow.get_sample_data_complete(feature)
targets = self.flow.get_target_data(is_binary, good_labels)
if len(samples) > 2:
algorithm.train(samples, targets)
pos_train = 0
neg_train = 0
for i in targets:
if i in good_labels:
neg_train += 1
else:
pos_train += 1
manager.add_new_result(pos_train, neg_train)
def default(self, algorithm, loader, feature, file):
print "Training size is " + str(loader.get_netflow().get_size()) + "."
if loader.get_netflow().get_size() <= 2:
print "Training set too small."
return None
#else:
# samples = loader.get_netflow().get_sample_data_complete(feature)
# targets = loader.get_netflow().get_target_data()
# #try:
# algorithm.train(samples, targets)
# #except Exception as e:
# # raise e
# # print "Wrong training data set used."
# # return False
print "Training set \"" + file + "\" done."
return loader.get_netflow()
