File size: 2,648 Bytes
62a2f1c |
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 |
from abc import ABCMeta, abstractmethod
import torch
import torch.nn as nn
from pytorch_lightning.utilities import rank_zero_warn
__all__ = ["Atomref"]
class BasePrior(nn.Module, metaclass=ABCMeta):
"""
Base class for prior models.
Derive this class to make custom prior models, which take some arguments and a dataset as input.
As an example, have a look at the `torchmdnet.priors.Atomref` prior.
"""
def __init__(self):
super(BasePrior, self).__init__()
@abstractmethod
def get_init_args(self):
"""
A function that returns all required arguments to construct a prior object.
The values should be returned inside a dict with the keys being the arguments' names.
All values should also be saveable in a .yaml file as this is used to reconstruct the
prior model from a checkpoint file.
"""
return
@abstractmethod
def forward(self, x, z):
"""
Forward method of the prior model.
Args:
x (torch.Tensor): scalar atomwise predictions from the model.
z (torch.Tensor): atom types of all atoms.
Returns:
torch.Tensor: updated scalar atomwise predictions
"""
return
class Atomref(BasePrior):
"""
Atomref prior model.
When using this in combination with some dataset, the dataset class must implement
the function `get_atomref`, which returns the atomic reference values as a tensor.
"""
def __init__(self, max_z=None, dataset=None):
super(Atomref, self).__init__()
if max_z is None and dataset is None:
raise ValueError("Can't instantiate Atomref prior, all arguments are None.")
if dataset is None:
atomref = torch.zeros(max_z, 1)
else:
atomref = dataset.get_atomref()
if atomref is None:
rank_zero_warn(
"The atomref returned by the dataset is None, defaulting to zeros with max. "
"atomic number 99. Maybe atomref is not defined for the current target."
)
atomref = torch.zeros(100, 1)
if atomref.ndim == 1:
atomref = atomref.view(-1, 1)
self.register_buffer("initial_atomref", atomref)
self.atomref = nn.Embedding(len(atomref), 1)
self.atomref.weight.data.copy_(atomref)
def reset_parameters(self):
self.atomref.weight.data.copy_(self.initial_atomref)
def get_init_args(self):
return dict(max_z=self.initial_atomref.size(0))
def forward(self, x, z):
return x + self.atomref(z)
|