model.py 1.16 KB
Newer Older
1
2
3
4
5
import torch
import torchani
import os


6
7
8
9
10
11
12
13
14
15
16
def atomic():
    model = torch.nn.Sequential(
        torch.nn.Linear(384, 128),
        torch.nn.CELU(0.1),
        torch.nn.Linear(128, 128),
        torch.nn.CELU(0.1),
        torch.nn.Linear(128, 64),
        torch.nn.CELU(0.1),
        torch.nn.Linear(64, 1)
    )
    return model
17
18
19


def get_or_create_model(filename, benchmark=False,
20
21
22
                        device=torch.device('cpu')):
    aev_computer = torchani.SortedAEV(benchmark=benchmark)
    prepare = torchani.PrepareInput(aev_computer.species)
23
24
25
26
    model = torchani.models.CustomModel(
        reducer=torch.sum,
        benchmark=benchmark,
        per_species={
27
28
29
30
            'C': atomic(),
            'H': atomic(),
            'N': atomic(),
            'O': atomic(),
31
        })
32
33
34
35

    class Flatten(torch.nn.Module):

        def forward(self, x):
36
            return x[0], x[1].flatten()
37
38

    model = torch.nn.Sequential(prepare, aev_computer, model, Flatten())
39
40
41
42
    if os.path.isfile(filename):
        model.load_state_dict(torch.load(filename))
    else:
        torch.save(model.state_dict(), filename)
Gao, Xiang's avatar
Gao, Xiang committed
43
    return model.to(device), torchani.EnergyShifter(aev_computer.species)