| import torch |
| import numpy as np |
| from graph_decoder.diffusion_model import GraphDiT |
|
|
| def count_parameters(model): |
| r""" |
| Returns the number of trainable parameters and number of all parameters in the model. |
| """ |
| trainable_params, all_param = 0, 0 |
| for param in model.parameters(): |
| num_params = param.numel() |
| all_param += num_params |
| if param.requires_grad: |
| trainable_params += num_params |
|
|
| return trainable_params, all_param |
|
|
| def load_graph_decoder(path='model_labeled'): |
| model_config_path = f"{path}/config.yaml" |
| data_info_path = f"{path}/data.meta.json" |
|
|
| model = GraphDiT( |
| model_config_path=model_config_path, |
| data_info_path=data_info_path, |
| |
| model_dtype=torch.float32, |
| ) |
| model.init_model(path) |
| model.disable_grads() |
|
|
| trainable_params, all_param = count_parameters(model) |
| param_stats = "Loaded Graph DiT from {} trainable params: {:,} || all params: {:,} || trainable%: {:.4f}".format( |
| path, trainable_params, all_param, 100 * trainable_params / all_param |
| ) |
| print(param_stats) |
| return model |
|
|
