File size: 5,853 Bytes
bf928ee | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 | #!/usr/bin/env python3
"""Evaluate the official test set and weights on a single DCU."""
import argparse
import pickle
import time
from pathlib import Path
from _bootstrap import DATA_DIR, WEIGHT_DIR
import numpy as np
import tensorflow as tf
import tf_keras
from sklearn.metrics import average_precision_score, roc_auc_score
DEFAULT_DATA_DIR = DATA_DIR / "features" / "cdr_kmer3_ag_kmer1"
DEFAULT_WEIGHT_PATH = WEIGHT_DIR / "cdr_kmer3_ag_kmer1" / "Model99.h5"
def parse_args():
parser = argparse.ArgumentParser()
parser.add_argument("--data-dir", type=Path, default=DEFAULT_DATA_DIR)
parser.add_argument("--weights", type=Path, default=DEFAULT_WEIGHT_PATH)
parser.add_argument("--batch-size", type=int, default=64)
parser.add_argument(
"--output",
type=Path,
help="optional .npz path for labels and predictions",
)
return parser.parse_args()
def load_test_set(data_dir):
# Register features.InputFeatures before loading the official pickles.
import model # noqa: F401
with (data_dir / "cdr_features_te.pickle").open("rb") as reader:
cdr_features = pickle.load(reader)
with (data_dir / "ag_features_te.pickle").open("rb") as reader:
ag_features = pickle.load(reader)
if not cdr_features:
raise ValueError("The test set is empty")
if len(cdr_features) != len(ag_features):
raise ValueError(
f"Unpaired test features: CDR={len(cdr_features)}, AG={len(ag_features)}"
)
cdr_labels = np.asarray([item.label_id for item in cdr_features], dtype=np.int32)
ag_labels = np.asarray([item.label_id for item in ag_features], dtype=np.int32)
if not np.array_equal(cdr_labels, ag_labels):
mismatch_count = int(np.count_nonzero(cdr_labels != ag_labels))
raise ValueError(f"CDR/AG labels differ for {mismatch_count} test samples")
if set(np.unique(cdr_labels).tolist()) != {0, 1}:
raise ValueError(f"Expected binary labels 0/1, got {np.unique(cdr_labels)}")
inputs = [
np.asarray([item.input_ids for item in cdr_features], dtype=np.int32),
np.asarray([item.cdr_number_ids for item in cdr_features], dtype=np.int32),
np.asarray([item.input_ids for item in ag_features], dtype=np.int32),
]
sample_count = len(cdr_features)
expected_shapes = ((sample_count, 24), (sample_count, 24), (sample_count, 2371))
actual_shapes = tuple(array.shape for array in inputs)
if actual_shapes != expected_shapes:
raise ValueError(f"Unexpected input shapes: {actual_shapes}, expected {expected_shapes}")
return inputs, cdr_labels
def main():
args = parse_args()
if args.batch_size < 1:
raise ValueError(f"--batch-size must be positive, got {args.batch_size}")
if not args.data_dir.is_dir():
raise FileNotFoundError(f"Data directory not found: {args.data_dir}")
if not args.weights.is_file():
raise FileNotFoundError(f"Weight file not found: {args.weights}")
gpus = tf.config.list_physical_devices("GPU")
if not gpus:
raise RuntimeError(
"TensorFlow cannot see a DCU. Set HIP_VISIBLE_DEVICES and "
"CUDA_VISIBLE_DEVICES before starting Python."
)
for gpu in gpus:
tf.config.experimental.set_memory_growth(gpu, True)
inputs, labels = load_test_set(args.data_dir)
# Import after configuring the accelerator to avoid initializing it early.
from model import get_model
with tf.device("/GPU:0"):
model = get_model()
model.load_weights(str(args.weights))
start_time = time.perf_counter()
predictions = model.predict(
inputs,
batch_size=args.batch_size,
verbose=1,
).reshape(-1)
elapsed_seconds = time.perf_counter() - start_time
if predictions.shape != labels.shape:
raise ValueError(
f"Prediction/label shape mismatch: {predictions.shape} vs {labels.shape}"
)
if not np.all(np.isfinite(predictions)):
raise FloatingPointError("Test predictions contain NaN or Inf")
if np.any((predictions < 0.0) | (predictions > 1.0)):
raise ValueError("Test predictions are outside [0, 1]")
auroc = float(roc_auc_score(labels, predictions))
aupr = float(average_precision_score(labels, predictions))
positive_count = int(labels.sum())
negative_count = int(labels.size - positive_count)
dcu_variable_count = sum("GPU:0" in variable.device for variable in model.variables)
print("\nFull official test-set evaluation")
print(f"TensorFlow: {tf.__version__}")
print(f"tf_keras: {tf_keras.__version__}")
print(f"DCU: {gpus[0]}")
print(f"Samples: {labels.size}")
print(f"Positive/negative: {positive_count}/{negative_count}")
print(f"Batch size: {args.batch_size}")
print(f"Variables placed on GPU:0: {dcu_variable_count}/{len(model.variables)}")
print(f"Prediction min/max/mean: {predictions.min():.10g} / "
f"{predictions.max():.10g} / {predictions.mean():.10g}")
print(f"Elapsed seconds: {elapsed_seconds:.3f}")
print(f"Throughput samples/s: {labels.size / elapsed_seconds:.3f}")
print(f"AUROC: {auroc:.10f}")
print(f"AUPR: {aupr:.10f}")
print("Paper reference AUROC/AUPR: 0.926 / 0.952")
if dcu_variable_count == 0:
raise RuntimeError("No model variables were placed on GPU:0")
if args.output is not None:
args.output.parent.mkdir(parents=True, exist_ok=True)
np.savez_compressed(
args.output,
labels=labels,
predictions=predictions,
auroc=np.asarray(auroc),
aupr=np.asarray(aupr),
)
print(f"Saved predictions: {args.output}")
print("Full test-set inference: PASS")
if __name__ == "__main__":
main()
|