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)