(args, prior_model=None, mean=None, std=None)
| 472 | VIS_MP_MAP = {'Node': ViS_MP_Vertex_Node, 'Edge': ViS_MP_Vertex_Edge, 'None': ViS_MP} |
| 473 | |
| 474 | def create_model(args, prior_model=None, mean=None, std=None): |
| 475 | visnet_args = dict( |
| 476 | lmax=args["lmax"], |
| 477 | vecnorm_type=args["vecnorm_type"], |
| 478 | trainable_vecnorm=args["trainable_vecnorm"], |
| 479 | num_heads=args["num_heads"], |
| 480 | num_layers=args["num_layers"], |
| 481 | hidden_channels=args["embedding_dimension"], |
| 482 | num_rbf=args["num_rbf"], |
| 483 | rbf_type=args["rbf_type"], |
| 484 | trainable_rbf=args["trainable_rbf"], |
| 485 | activation=args["activation"], |
| 486 | attn_activation=args["attn_activation"], |
| 487 | max_z=args["max_z"], |
| 488 | cutoff=args["cutoff"], |
| 489 | max_num_neighbors=args["max_num_neighbors"], |
| 490 | vertex_type=args["vertex_type"], |
| 491 | ) |
| 492 | |
| 493 | # representation network |
| 494 | if args["model"] == "ViSNetBlock": |
| 495 | representation_model = ViSNetBlock(**visnet_args) |
| 496 | else: |
| 497 | raise ValueError(f"Unknown model {args['model']}.") |
| 498 | |
| 499 | # prior model |
| 500 | if args["prior_model"] and prior_model is None: |
| 501 | assert "prior_args" in args, ( |
| 502 | f"Requested prior model {args['prior_model']} but the " |
| 503 | f'arguments are lacking the key "prior_args".' |
| 504 | ) |
| 505 | assert hasattr(priors, args["prior_model"]), ( |
| 506 | f'Unknown prior model {args["prior_model"]}. ' |
| 507 | f'Available models are {", ".join(priors.__all__)}' |
| 508 | ) |
| 509 | # instantiate prior model if it was not passed to create_model (i.e. when loading a model) |
| 510 | prior_model = getattr(priors, args["prior_model"])(**args["prior_args"]) |
| 511 | |
| 512 | # create output network |
| 513 | output_prefix = "Equivariant" |
| 514 | output_model = getattr(output_modules, output_prefix + args["output_model"])(args["embedding_dimension"], args["activation"]) |
| 515 | |
| 516 | model = ViSNet( |
| 517 | representation_model, |
| 518 | output_model, |
| 519 | prior_model=prior_model, |
| 520 | reduce_op=args["reduce_op"], |
| 521 | mean=mean, |
| 522 | std=std, |
| 523 | derivative=args["derivative"], |
| 524 | ) |
| 525 | return model |
| 526 | |
| 527 | |
| 528 | def load_model(filepath, args=None, device="cpu", **kwargs): |
no test coverage detected