nnp_training.py 4.92 KB
Newer Older
1
2
3
4
import torch
import ignite
import torchani
import model
5
6
7
8
import tqdm
import timeit
import tensorboardX
import math
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
import argparse
import json

# parse command line arguments
parser = argparse.ArgumentParser()
parser.add_argument('dataset_path',
                    help='Path of the dataset, can a hdf5 file \
                          or a directory containing hdf5 files')
parser.add_argument('--dataset_checkpoint',
                    help='Checkpoint file for datasets',
                    default='dataset-checkpoint.dat')
parser.add_argument('--model_checkpoint',
                    help='Checkpoint file for model',
                    default='model.pt')
parser.add_argument('-m', '--max_epochs',
                    help='Maximum number of epoches',
25
26
27
28
                    default=100, type=int)
parser.add_argument('--training_rmse_every',
                    help='Compute training RMSE every epoches',
                    default=20, type=int)
29
30
31
parser.add_argument('-d', '--device',
                    help='Device of modules and tensors',
                    default=('cuda' if torch.cuda.is_available() else 'cpu'))
32
33
34
parser.add_argument('--batch_size',
                    help='Number of conformations of each batch',
                    default=1024, type=int)
35
36
37
38
39
40
41
42
43
parser.add_argument('--log',
                    help='Log directory for tensorboardX',
                    default=None)
parser.add_argument('--optimizer',
                    help='Optimizer used to train the model',
                    default='Adam')
parser.add_argument('--optim_args',
                    help='Arguments to optimizers, in the format of json',
                    default='{}')
44
45
46
parser.add_argument('--early_stopping',
                    help='Stop after epoches of no improvements',
                    default=math.inf, type=int)
47
48
49
50
51
parser = parser.parse_args()

# set up the training
device = torch.device(parser.device)
writer = tensorboardX.SummaryWriter(log_dir=parser.log)
52
start = timeit.default_timer()
53

Gao, Xiang's avatar
Gao, Xiang committed
54
55
nnp = model.get_or_create_model(parser.model_checkpoint, device=device)
shift_energy = torchani.buildins.energy_shifter
56
training, validation, testing = torchani.training.load_or_create(
57
    parser.dataset_checkpoint, parser.batch_size, model.consts.species,
58
59
60
    parser.dataset_path, device=device,
    transform=[shift_energy.subtract_from_dataset])
container = torchani.training.Container({'energies': nnp})
61
62
63
64

parser.optim_args = json.loads(parser.optim_args)
optimizer = getattr(torch.optim, parser.optimizer)
optimizer = optimizer(nnp.parameters(), **parser.optim_args)
65

66
trainer = ignite.engine.create_supervised_trainer(
67
    container, optimizer, torchani.training.MSELoss('energies'))
68
evaluator = ignite.engine.create_supervised_evaluator(container, metrics={
69
        'RMSE': torchani.training.RMSEMetric('energies')
70
71
72
    })


73
74
75
76
def hartree2kcal(x):
    return 627.509 * x


77
78
79
@trainer.on(ignite.engine.Events.STARTED)
def initialize(trainer):
    trainer.state.best_validation_rmse = math.inf
80
    trainer.state.no_improve_count = 0
81
82


83
84
85
86
87
@trainer.on(ignite.engine.Events.EPOCH_STARTED)
def init_tqdm(trainer):
    trainer.state.tqdm = tqdm.tqdm(total=len(training), desc='epoch')


88
@trainer.on(ignite.engine.Events.ITERATION_COMPLETED)
89
90
def update_tqdm(trainer):
    trainer.state.tqdm.update(1)
91
92
93


@trainer.on(ignite.engine.Events.EPOCH_COMPLETED)
94
95
def finalize_tqdm(trainer):
    trainer.state.tqdm.close()
96
97


98
@trainer.on(ignite.engine.Events.EPOCH_STARTED)
99
def validation_and_checkpoint(trainer):
100
    # compute validation RMSE
101
102
    evaluator.run(validation)
    metrics = evaluator.state.metrics
103
104
    rmse = hartree2kcal(metrics['RMSE'])
    writer.add_scalar('validation_rmse_vs_epoch', rmse, trainer.state.epoch)
105
106
107
108
109
110
111
112
113
114

    # compute training RMSE
    if trainer.state.epoch % parser.training_rmse_every == 0:
        evaluator.run(training)
        metrics = evaluator.state.metrics
        rmse = hartree2kcal(metrics['RMSE'])
        writer.add_scalar('training_rmse_vs_epoch', rmse,
                          trainer.state.epoch)

    # handle best validation RMSE
115
    if rmse < trainer.state.best_validation_rmse:
116
        trainer.state.no_improve_count = 0
117
        trainer.state.best_validation_rmse = rmse
118
119
        writer.add_scalar('best_validation_rmse_vs_epoch', rmse,
                          trainer.state.epoch)
120
        torch.save(nnp.state_dict(), parser.model_checkpoint)
121
122
123
124
125
    else:
        trainer.state.no_improve_count += 1

    if trainer.state.no_improve_count > parser.early_stopping:
        trainer.terminate()
126
127
128
129
130
131
132
133
134
135
136
137


@trainer.on(ignite.engine.Events.EPOCH_STARTED)
def log_time(trainer):
    elapsed = round(timeit.default_timer() - start, 2)
    writer.add_scalar('time_vs_epoch', elapsed, trainer.state.epoch)


@trainer.on(ignite.engine.Events.ITERATION_COMPLETED)
def log_loss_and_time(trainer):
    iteration = trainer.state.iteration
    rmse = hartree2kcal(math.sqrt(trainer.state.output))
138
    writer.add_scalar('training_atomic_rmse_vs_iteration', rmse, iteration)
139
140


141
trainer.run(training, max_epochs=parser.max_epochs)