def test_predict():
dd = DGADetector()
test_domains = cudf.Series(["nvidia.com", "dfsdfsdf"])
dd.load_model(model_filepath)
actual_output = dd.predict(test_domains)
expected_output = cudf.Series([1, 0])
assert actual_output.equals(actual_output)
def test_load_model():
dd = DGADetector()
dd.load_model(model_filepath)
gpu_count = torch.cuda.device_count()
if gpu_count > 1:
assert isinstance(dd.model, nn.DataParallel)
else:
assert isinstance(dd.model, RNNClassifier)
def test_predict():
if torch.cuda.is_available():
dd = DGADetector()
test_domains = cudf.Series(["nvidia.com", "dfsdfsdf"])
dd.load_model(MODEL_FILENAME)
actual_output = dd.predict(test_domains)
expected_output = cudf.Series([1, 0])
assert actual_output.equals(expected_output)
def worker_init():
# Initialization for each dask worker
from clx.analytics.dga_detector import DGADetector
worker = dask.distributed.get_worker()
dd = DGADetector()
print("Initializing Dask worker: " + str(worker) +
" with dga model. Model File: " + str(args.model))
dd.load_model(args.model)
worker.data["dga_detector"] = dd
print("Successfully initialized dask worker " + str(worker))
def worker_init(self):
# Initialization for each dask worker
from clx.analytics.dga_detector import DGADetector
worker = dask.distributed.get_worker()
dd = DGADetector()
print("Initializing Dask worker: " + str(worker) +
" with dga model. Model File: " + str(self.args.model))
dd.load_model(self.args.model)
# this dict can be used for adding more objects to distributed dask worker
obj_dict = {"dga_detector": dd}
worker = utils.init_dask_workers(worker, self.config, obj_dict)
def test_load_model(tmpdir):
if torch.cuda.is_available():
# save model
dd.save_model(str(tmpdir.join("clx_dga.mdl")))
assert path.exists(str(tmpdir.join("clx_dga.mdl")))
# load model
dd2 = DGADetector()
dd2.init_model()
dd2.load_model(str(tmpdir.join("clx_dga.mdl")))
gpu_count = torch.cuda.device_count()
if gpu_count > 1:
assert isinstance(dd2.model.module, RNNClassifier)
else:
assert isinstance(dd2.model, RNNClassifier)
def test_load_model():
if torch.cuda.is_available():
dd = DGADetector()
dd.load_model(MODEL_FILENAME)
assert isinstance(dd.model, RNNClassifier)