Spaces:
Sleeping
Sleeping
| import torch | |
| import torch.nn.functional as F | |
| from torch import Tensor | |
| from torch.nn import ( | |
| Sequential, | |
| Linear, | |
| ReLU, | |
| GRU, | |
| Embedding, | |
| BatchNorm1d, | |
| Dropout, | |
| LayerNorm, | |
| ) | |
| from torch_geometric.nn import ( | |
| Set2Set, | |
| global_mean_pool, | |
| global_add_pool, | |
| global_max_pool, | |
| ) | |
| # Sine matrix with neural network | |
| class SM(torch.nn.Module): | |
| def __init__(self, data, dim1=64, fc_count=1, **kwargs): | |
| super(SM, self).__init__() | |
| self.lin1 = torch.nn.Linear(data[0].extra_features_SM.shape[1], dim1) | |
| self.lin_list = torch.nn.ModuleList( | |
| [torch.nn.Linear(dim1, dim1) for i in range(fc_count)] | |
| ) | |
| self.lin2 = torch.nn.Linear(dim1, 1) | |
| def forward(self, data): | |
| out = F.relu(self.lin1(data.extra_features_SM)) | |
| for layer in self.lin_list: | |
| out = F.relu(layer(out)) | |
| out = self.lin2(out) | |
| if out.shape[1] == 1: | |
| return out.view(-1) | |
| else: | |
| return out | |
| # Smooth Overlap of Atomic Positions with neural network | |
| class SOAP(torch.nn.Module): | |
| def __init__(self, data, dim1, fc_count, **kwargs): | |
| super(SOAP, self).__init__() | |
| self.lin1 = torch.nn.Linear(data[0].extra_features_SOAP.shape[1], dim1) | |
| self.lin_list = torch.nn.ModuleList( | |
| [torch.nn.Linear(dim1, dim1) for i in range(fc_count)] | |
| ) | |
| self.lin2 = torch.nn.Linear(dim1, 1) | |
| def forward(self, data): | |
| out = F.relu(self.lin1(data.extra_features_SOAP)) | |
| for layer in self.lin_list: | |
| out = F.relu(layer(out)) | |
| out = self.lin2(out) | |
| if out.shape[1] == 1: | |
| return out.view(-1) | |
| else: | |
| return out | |