Spaces:
Sleeping
Sleeping
File size: 1,817 Bytes
dfc4f2b | 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 | 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
|