| |
| """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() |
|
|
