File size: 2,929 Bytes
3e02ab8 | 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 | import torch
from typing import Union, Sequence, Dict
from math import sqrt
from onescience.datapipes.materials.nequip import AtomicDataDict
class AvgNumNeighborsNorm(torch.nn.Module):
def __init__(
self,
type_names: Sequence[str],
avg_num_neighbors: Union[float, Dict[str, float]],
) -> None:
"""
Module to normalize features during training using per type edge sum normalization.
Args:
type_names (Sequence[str]): list of atom type names
avg_num_neighbors (float/Dict[str, float]): used to normalize edge sums for better numerics
"""
super().__init__()
assert avg_num_neighbors is not None, "avg_num_neighbors must be specified"
self.in_field = self.out_field = AtomicDataDict.NODE_FEATURES_KEY
self.norm_key = AtomicDataDict.FEATURE_NORM_FACTOR_KEY
# Put avg_num_neighbors in a list (global or per type)
if isinstance(avg_num_neighbors, (float, int)):
avg_num_neighbors = [avg_num_neighbors]
elif isinstance(avg_num_neighbors, dict):
assert set(type_names) == set(avg_num_neighbors.keys())
avg_num_neighbors = [avg_num_neighbors[k] for k in type_names]
else:
raise RuntimeError(
"Unrecognized format for `avg_num_neighbors`, only floats or dicts allowed."
)
assert isinstance(avg_num_neighbors, list)
# Tensorize avg_num_neighbors and register as buffer
norm_const = torch.tensor([(1.0 / sqrt(N)) for N in avg_num_neighbors])
norm_const = norm_const.reshape(-1, 1)
# Persistent=False to ensure backwards compatibility of FMs.
# TODO remove this once we're sure FMs are not using this anymore
self.register_buffer("norm_const", norm_const, persistent=False)
# If global avg_num_neighbors or only one type, no need to do embedding lookup in forward
self.norm_shortcut = self.norm_const.numel() == 1
def forward(self, data: AtomicDataDict.Type) -> AtomicDataDict.Type:
features = data[self.in_field]
norm_size = features.size(0)
if self.norm_key in data and data[self.norm_key].size(0) == norm_size:
norm_factor = data[self.norm_key]
else:
# Compute norm factor for the first time
if self.norm_shortcut:
# No need to do embedding lookup in forward
norm_factor = self.norm_const.expand(norm_size, -1)
else:
# Embed each avg_num_neighbors value per type
norm_factor = torch.nn.functional.embedding(
data[AtomicDataDict.ATOM_TYPE_KEY][:norm_size],
self.norm_const,
)
data[self.norm_key] = norm_factor # shape: (num_local_nodes, 1)
data[self.out_field] = norm_factor * features
return data
|