MCPcopy Create free account

hub / github.com/atomicarchitects/equiformer_v2 / functions

Functions244 in github.com/atomicarchitects/equiformer_v2

↓ 34 callersMethodinfo
(self, *args)
oc20/trainer/logger.py:49
↓ 27 callersMethodupdate
(self, val, n=1)
oc20/trainer/engine.py:23
↓ 12 callersMethodset_embedding
(self, embedding)
nets/equiformer_v2/so3.py:244
↓ 12 callersMethodstep
(self, metrics=None, epoch=None)
oc20/trainer/lr_scheduler.py:147
↓ 11 callersMethodsave
( self, metrics=None, checkpoint_file="checkpoint.pt", training_state=True,
oc20/trainer/base_trainer_oc20.py:511
↓ 10 callersMethodset_lmax_mmax
(self, lmax_list, mmax_list)
nets/equiformer_v2/so3.py:250
↓ 9 callersMethodcoefficient_idx
(self, lmax, mmax)
nets/equiformer_v2/so3.py:139
↓ 9 callersMethodupdate
(self, val, n=1)
engine.py:23
↓ 8 callersMethod__init__
(self, in_channels, out_channels, bias=True)
nets/equiformer_v2/activation.py:47
↓ 8 callersMethodinfo
(self, *args)
logger.py:49
↓ 6 callersMethodclose_db
(self)
oc20/trainer/lmdb_dataset.py:134
↓ 5 callersMethod__init__
( self, lmax, mmax, normalization='integral', resolution=None, )
nets/equiformer_v2/so3.py:493
↓ 5 callersMethodload
(self)
oc20/trainer/base_trainer_v2.py:308
↓ 4 callersMethod__init__
(self, drop_prob=None)
nets/equiformer_v2/drop.py:34
↓ 4 callersMethod__init__
(self, lmax, num_channels, eps=1e-5, affine=True, normalization='component')
nets/equiformer_v2/layer_norm.py:37
↓ 4 callersMethod_compute_metrics
(self, out, batch_list, evaluator, metrics={})
oc20/trainer/forces_trainer_v2.py:575
↓ 4 callersMethodclone
(self)
nets/equiformer_v2/so3.py:231
↓ 4 callersMethodget_dataloader
(self, dataset, sampler)
oc20/trainer/base_trainer_oc20.py:265
↓ 4 callersMethodget_sampler
(self, dataset, batch_size, shuffle)
oc20/trainer/base_trainer_oc20.py:245
↓ 4 callersMethodtrain
Derived classes should implement this function.
oc20/trainer/base_trainer_oc20.py:608
↓ 3 callersMethod_forward
(self, batch_list)
oc20/trainer/forces_trainer_v2.py:447
↓ 3 callersMethod_l_primary
(self, mapping)
nets/equiformer_v2/so3.py:293
↓ 3 callersFunction_z_rot_mat
(angle, l)
nets/equiformer_v2/wigner.py:30
↓ 3 callersMethodget_from_grid_mat
(self, device)
nets/equiformer_v2/so3.py:563
↓ 3 callersFunctionget_normalization_layer
(norm_type, lmax, num_channels, eps=1e-5, affine=True, normalization='component')
nets/equiformer_v2/layer_norm.py:13
↓ 3 callersMethodget_to_grid_mat
(self, device)
nets/equiformer_v2/so3.py:558
↓ 3 callersMethodload
(self)
oc20/trainer/base_trainer_oc20.py:209
↓ 2 callersMethod__init__
(self, config, transform=None)
oc20/trainer/lmdb_dataset.py:59
↓ 2 callersMethod__init__
( self, sphere_channels, m_output_channels, lmax_list, mmax_list,
nets/equiformer_v2/so2_ops.py:220
↓ 2 callersMethod__init__
( self, sphere_channels, attn_hidden_channels, num_heads, attn_alpha_c
nets/equiformer_v2/transformer_block.py:494
↓ 2 callersMethod_backward
(self, loss)
oc20/trainer/base_trainer_v2.py:546
↓ 2 callersMethod_compute_loss
(self, out, batch_list)
oc20/trainer/forces_trainer_v2.py:468
↓ 2 callersMethod_compute_loss
Derived classes should implement this function.
oc20/trainer/base_trainer_oc20.py:683
↓ 2 callersMethod_expand_edge
(self, edge_index)
nets/equiformer_v2/so3.py:256
↓ 2 callersMethod_forward
(self, batch_list)
oc20/trainer/energy_trainer_v2.py:398
↓ 2 callersMethod_forward
Derived classes should implement this function.
oc20/trainer/base_trainer_oc20.py:679
↓ 2 callersMethod_m_primary
(self, mapping)
nets/equiformer_v2/so3.py:288
↓ 2 callersFunction_mask_input
(inputs: torch.Tensor, mask: torch.Tensor)
oc20/trainer/energy_trainer_v2.py:30
↓ 2 callersMethod_reduce_edge
(self, edge_index, num_nodes)
nets/equiformer_v2/so3.py:275
↓ 2 callersMethod_rotate_inv
(self, SO3_rotation, mappingReduced)
nets/equiformer_v2/so3.py:319
↓ 2 callersMethodconnect_db
(self, lmdb_path=None)
oc20/trainer/lmdb_dataset.py:122
↓ 2 callersFunctiondrop_path
Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks). This is the same as the DropConnect impl I created for Ef
nets/equiformer_v2/drop.py:13
↓ 2 callersFunctionget_l_to_all_m_expand_index
(lmax)
nets/equiformer_v2/layer_norm.py:26
↓ 2 callersMethodget_lr
(self)
oc20/trainer/lr_scheduler.py:175
↓ 2 callersMethodhpo_update
( self, epoch, step, train_metrics, val_metrics, test_metrics=None )
oc20/trainer/base_trainer_oc20.py:581
↓ 2 callersFunctioninterpolate_init_relaxed_pos
(batch)
oc20/trainer/base_trainer_v2.py:81
↓ 2 callersFunctionis_dist_avail_and_initialized
()
utils.py:17
↓ 2 callersMethodload_datasets
(self)
oc20/trainer/base_trainer_oc20.py:275
↓ 2 callersMethodload_logger
(self)
oc20/trainer/base_trainer_oc20.py:232
↓ 2 callersMethodload_loss
(self)
oc20/trainer/base_trainer_oc20.py:443
↓ 2 callersMethodload_task
Initialize task-specific information. Derived classes should implement this function.
oc20/trainer/base_trainer_oc20.py:348
↓ 2 callersMethodno_weight_decay
(self)
nets/equiformer_v2/equiformer_v2_oc20.py:512
↓ 2 callersFunctionpyg2_data_transform
(data: Data)
oc20/trainer/lmdb_dataset.py:34
↓ 2 callersMethodrotate
(self, embedding, out_lmax, out_mmax)
nets/equiformer_v2/so3.py:445
↓ 2 callersMethodrotate_inv
(self, embedding, in_lmax, in_mmax)
nets/equiformer_v2/so3.py:452
↓ 2 callersMethodrun_relaxations
(self, split="val")
oc20/trainer/forces_trainer_v2.py:619
↓ 2 callersMethodsave_results
(self, predictions, results_file, keys)
oc20/trainer/base_trainer_oc20.py:721
↓ 2 callersMethodupdate_best
( self, primary_metric, val_metrics, disable_eval_tqdm=True, )
oc20/trainer/forces_trainer_v2.py:269
↓ 2 callersMethodvalidate
(self, split="val", disable_tqdm=False, use_ema=False)
oc20/trainer/forces_trainer_v2.py:848
↓ 2 callersMethodvalidate
(self, split="val", disable_tqdm=False, use_ema=True)
oc20/trainer/base_trainer_v2.py:478
↓ 1 callersMethodRotationToWignerDMatrix
(self, edge_rot_mat, start_lmax, end_lmax)
nets/equiformer_v2/so3.py:461
↓ 1 callersMethod_compute_auxiliary_task_weight
(self)
oc20/trainer/energy_trainer_v2.py:462
↓ 1 callersMethod_compute_loss
(self, out, batch_list)
oc20/trainer/energy_trainer_v2.py:413
↓ 1 callersMethod_compute_metrics
(self, out, batch_list, evaluator, metrics={})
oc20/trainer/energy_trainer_v2.py:445
↓ 1 callersMethod_forward_otf_graph
(self, data)
oc20/trainer/base_trainer_v2.py:651
↓ 1 callersMethod_forward_use_pbc
(self, data)
oc20/trainer/base_trainer_v2.py:664
↓ 1 callersMethod_from_grid
(self, x_grid, SO3_grid, lmax=-1)
nets/equiformer_v2/so3.py:391
↓ 1 callersMethod_init_edge_rot_mat
(self, data, edge_index, edge_distance_vec)
nets/equiformer_v2/equiformer_v2_oc20.py:474
↓ 1 callersMethod_mask_input
(self, inputs, mask)
oc20/trainer/base_trainer_v2.py:688
↓ 1 callersMethod_rotate
(self, SO3_rotation, lmax_list, mmax_list)
nets/equiformer_v2/so3.py:298
↓ 1 callersFunctionadd_weight_decay
(model, weight_decay=1e-5, skip_list=())
optim_factory.py:27
↓ 1 callersFunctionadd_weight_decay
(model, weight_decay, skip_list=())
oc20/trainer/base_trainer_v2.py:60
↓ 1 callersMethodcomplex_idx
Add `m_complex` and `l_harmonic` to the input arguments since we cannot use `self.m_complex`.
nets/equiformer_v2/so3.py:112
↓ 1 callersMethodcompute_stats
Compute mean of numbers of nodes and edges Assume using cpu
oc20/trainer/base_trainer_v2.py:594
↓ 1 callersFunctioncreate_optimizer_v2
Create an optimizer. TODO currently the model is passed in and all parameters are selected for optimization. For more general use an interfa
optim_factory.py:74
↓ 1 callersMethodfilter_kwargs
(self, config)
oc20/trainer/lr_scheduler.py:160
↓ 1 callersFunctiongaussian
(x, mean, std)
nets/equiformer_v2/gaussian_rbf.py:5
↓ 1 callersMethodget_logger
(self, output_dir, log_to_file)
logger.py:17
↓ 1 callersMethodget_logger
(self, output_dir, log_to_file)
oc20/trainer/logger.py:17
↓ 1 callersFunctionget_rank
()
utils.py:31
↓ 1 callersMethodget_rotate_inv_rescale
(self, lmax, mmax)
nets/equiformer_v2/so3.py:159
↓ 1 callersFunctioninit_edge_rot_mat
(edge_distance_vec)
nets/equiformer_v2/edge_rot_mat.py:4
↓ 1 callersFunctionis_main_process
()
utils.py:37
↓ 1 callersMethodload_extras
(self)
oc20/trainer/base_trainer_v2.py:441
↓ 1 callersMethodload_extras
(self)
oc20/trainer/base_trainer_oc20.py:499
↓ 1 callersMethodload_model
(self)
oc20/trainer/base_trainer_v2.py:333
↓ 1 callersMethodload_model
(self)
oc20/trainer/base_trainer_oc20.py:351
↓ 1 callersMethodload_optimizer
(self)
oc20/trainer/base_trainer_v2.py:387
↓ 1 callersMethodload_optimizer
(self)
oc20/trainer/base_trainer_oc20.py:461
↓ 1 callersMethodload_seed_from_config
(self)
oc20/trainer/base_trainer_v2.py:319
↓ 1 callersMethodload_seed_from_config
(self)
oc20/trainer/base_trainer_oc20.py:219
↓ 1 callersFunctionmain
(args)
oc20/trainer/make_lmdb_sizes.py:34
↓ 1 callersFunctionmultiply
(obj, num)
oc20/trainer/lr_scheduler.py:11
↓ 1 callersFunctionoptimizer_kwargs
cfg/argparse to kwargs helper Convert optimizer args in argparse args or cfg like object to keyword args for updated create fn.
optim_factory.py:45
↓ 1 callersMethodpredict
( self, data_loader, per_image=True, results_file=None, disable_tqdm=F
oc20/trainer/forces_trainer_v2.py:160
↓ 1 callersMethodpredict
( self, loader, per_image=True, results_file=None, disable_tqdm=False )
oc20/trainer/energy_trainer_v2.py:134
↓ 1 callersMethodreset
(self)
engine.py:17
↓ 1 callersMethodreset
(self)
oc20/trainer/engine.py:17
↓ 1 callersMethodrun
(self)
oc20/trainer/task_relaxation.py:7
↓ 1 callersMethodsave_hpo
(self, epoch, step, metrics, checkpoint_every)
oc20/trainer/base_trainer_oc20.py:570
next →1–100 of 244, ranked by callers