def test_train(shared_artifact_dir, cfg):
records_train = gen_records(NUM_SAMPLES)
records_validation = gen_records(NUM_SAMPLES)
bm = engine.BarrageModel(shared_artifact_dir)
net = bm.train(cfg, records_train, records_validation)
assert isinstance(net, tf.keras.models.Model)
assert os.path.isdir(shared_artifact_dir)
assert os.path.isdir(os.path.join(shared_artifact_dir, "dataset"))
assert os.path.isdir(os.path.join(shared_artifact_dir, "best_checkpoint"))
assert os.path.isdir(
os.path.join(shared_artifact_dir, "resume_checkpoints"))
assert os.path.isdir(os.path.join(shared_artifact_dir, "TensorBoard"))
assert os.path.isfile(
os.path.join(shared_artifact_dir, "training_report.csv"))
assert os.path.isfile(os.path.join(shared_artifact_dir, "config.json"))
assert os.path.isfile(os.path.join(shared_artifact_dir, "config.pkl"))
assert os.path.isfile(
os.path.join(shared_artifact_dir, "network_params.json"))
assert os.path.isfile(
os.path.join(shared_artifact_dir, "network_params.pkl"))
assert os.path.isfile(
os.path.join(shared_artifact_dir, "best_checkpoint",
"model_best.ckpt.index"))
assert os.path.isfile(
os.path.join(shared_artifact_dir, "resume_checkpoints",
"model_epoch_0001.ckpt.index"))
assert os.path.isfile(
os.path.join(shared_artifact_dir, "resume_checkpoints",
"model_epoch_0002.ckpt.index"))
def test_predict(shared_artifact_dir, records_score):
bm = engine.BarrageModel(shared_artifact_dir)
bm.load()
scores = bm.predict(records_score)
assert len(scores) == NUM_SAMPLES
for s in scores:
assert list(s.keys()) == ["output"]
assert len(s["output"]) == 2
assert 98 <= s["output"][0] <= 102
assert 98 <= s["output"][1] <= 102