File size: 3,706 Bytes
d46a58d | 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 | import logging
import os
import sys
from pathlib import Path
import numpy as np
import torch
PROJECT_ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(PROJECT_ROOT))
from model.meshgraphnet import MeshGraphNet
from onescience.utils.YParams import YParams
from onescience.launch.utils import load_checkpoint
from fake_data import build_cylinder_flow_datapipe
def build_model(model_params, device):
mlp_act = "silu" if model_params.recompute_activation else "relu"
return MeshGraphNet(
input_dim_nodes=model_params.num_input_features,
input_dim_edges=model_params.num_edge_features,
output_dim=model_params.num_output_features,
processor_size=model_params.processor_size,
hidden_dim_processor=model_params.hidden_dim_processor,
num_layers_node_processor=model_params.num_layers_node_processor,
num_layers_edge_processor=model_params.num_layers_edge_processor,
hidden_dim_node_encoder=model_params.hidden_dim_node_encoder,
hidden_dim_edge_encoder=model_params.hidden_dim_edge_encoder,
hidden_dim_node_decoder=model_params.hidden_dim_node_decoder,
mlp_activation_fn=mlp_act,
do_concat_trick=model_params.do_concat_trick,
num_processor_checkpoint_segments=model_params.num_processor_checkpoint_segments,
recompute_activation=model_params.recompute_activation,
).to(device)
def resolve_device(device_name: str):
if device_name == "cpu":
return torch.device("cpu")
if device_name in ("cuda", "gpu"):
if not torch.cuda.is_available():
raise RuntimeError("Config requested cuda device, but torch.cuda.is_available() is false.")
return torch.device("cuda:0")
return torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
def main():
os.chdir(PROJECT_ROOT)
logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
logger = logging.getLogger("mesh_graph_net.inference")
config_path = PROJECT_ROOT / "config" / "config.yaml"
cfg_model = YParams(config_path, "model")
cfg_data = YParams(config_path, "datapipe")
cfg_train = YParams(config_path, "training")
cfg_inference = YParams(config_path, "inference")
model_params = cfg_model.specific_params[cfg_model.name]
device = resolve_device(getattr(cfg_inference, "device", "auto"))
logger.info("Using device: %s", device)
datapipe = build_cylinder_flow_datapipe(
params=cfg_data,
distributed=False,
project_root=PROJECT_ROOT,
)
loader = datapipe.test_dataloader()
model = build_model(model_params, device)
checkpoint_dir = PROJECT_ROOT / getattr(cfg_inference, "checkpoint_dir", cfg_train.checkpoint_dir)
epoch = load_checkpoint(checkpoint_dir, models=model, device=device)
if epoch == 0:
logger.warning("No checkpoint found in %s; running with randomly initialized weights", checkpoint_dir)
model.eval()
predictions, targets = [], []
with torch.no_grad():
for batch in loader:
graph = batch[0] if isinstance(batch, (tuple, list)) else batch
graph = graph.to(device)
pred = model(graph.ndata["x"], graph.edata["x"], graph)
predictions.append(pred.cpu().numpy())
targets.append(graph.ndata["y"].cpu().numpy())
output_path = PROJECT_ROOT / cfg_inference.output_path
output_path.parent.mkdir(parents=True, exist_ok=True)
np.savez(output_path, prediction=np.concatenate(predictions, axis=0), target=np.concatenate(targets, axis=0))
logger.info("Saved inference results to %s", output_path)
if __name__ == "__main__":
main()
|