File size: 1,608 Bytes
986e0b8 | 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 | #!/usr/bin/env python3
"""Check the OneScience/DTK TensorFlow runtime and build Saluki."""
from __future__ import annotations
import argparse
import json
from pathlib import Path
from _bootstrap import configure_runtime
ROOT = configure_runtime()
def main() -> None:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--params", type=Path, default=ROOT / "conf" / "params.json")
parser.add_argument("--forward", action="store_true", help="Run one zero-input forward pass.")
args = parser.parse_args()
import numpy as np
import tensorflow as tf
from model.saluki import SalukiModel, load_params
params = load_params(args.params)
model = SalukiModel(params, head=0)
report = {
"tensorflow": tf.__version__,
"built_with_gpu": tf.test.is_built_with_gpu_support(),
"gpus": [device.name for device in tf.config.list_physical_devices("GPU")],
"input_shape": list(model.keras_model.input_shape),
"output_shape": list(model.keras_model.output_shape),
"parameters": int(model.keras_model.count_params())
}
if args.forward:
shape = model.keras_model.input_shape
output = model.keras_model(
np.zeros((1, shape[1], shape[2]), dtype=np.float32), training=False
)
output_np = output.numpy()
report["forward"] = {
"shape": list(output_np.shape),
"finite": bool(np.isfinite(output_np).all()),
"device": output.device
}
print(json.dumps(report, indent=2))
if __name__ == "__main__":
main()
|