def get_complete_network(aev_computer, device, return_networks=False):
aev_dim = aev_computer.aev_length
H_network = torch.nn.Sequential(torch.nn.Linear(aev_dim, 160),
torch.nn.CELU(0.1),
torch.nn.Linear(160,
128), torch.nn.CELU(0.1),
torch.nn.Linear(128, 96),
torch.nn.CELU(0.1), torch.nn.Linear(96, 1))
C_network = torch.nn.Sequential(torch.nn.Linear(aev_dim, 144),
torch.nn.CELU(0.1),
torch.nn.Linear(144,
112), torch.nn.CELU(0.1),
torch.nn.Linear(112, 96),
torch.nn.CELU(0.1), torch.nn.Linear(96, 1))
N_network = torch.nn.Sequential(torch.nn.Linear(aev_dim, 128),
torch.nn.CELU(0.1),
torch.nn.Linear(128,
112), torch.nn.CELU(0.1),
torch.nn.Linear(112, 96),
torch.nn.CELU(0.1), torch.nn.Linear(96, 1))
O_network = torch.nn.Sequential(torch.nn.Linear(aev_dim, 128),
torch.nn.CELU(0.1),
torch.nn.Linear(128,
112), torch.nn.CELU(0.1),
torch.nn.Linear(112, 96),
torch.nn.CELU(0.1), torch.nn.Linear(96, 1))
nn = torchani.ANIModel([H_network, C_network, N_network, O_network])
model = torchani.nn.Sequential(aev_computer, nn).to(device)
if return_networks:
return H_network, C_network, N_network, O_network, model, nn
return model, nn
C_network = torch.nn.Sequential(torch.nn.Linear(aev_dim, 144),
torch.nn.CELU(0.1), torch.nn.Linear(144, 112),
torch.nn.CELU(0.1), torch.nn.Linear(112, 96),
torch.nn.CELU(0.1), torch.nn.Linear(96, 1))
N_network = torch.nn.Sequential(torch.nn.Linear(aev_dim, 128),
torch.nn.CELU(0.1), torch.nn.Linear(128, 112),
torch.nn.CELU(0.1), torch.nn.Linear(112, 96),
torch.nn.CELU(0.1), torch.nn.Linear(96, 1))
O_network = torch.nn.Sequential(torch.nn.Linear(aev_dim, 128),
torch.nn.CELU(0.1), torch.nn.Linear(128, 112),
torch.nn.CELU(0.1), torch.nn.Linear(112, 96),
torch.nn.CELU(0.1), torch.nn.Linear(96, 1))
nn = torchani.ANIModel([H_network, C_network, N_network, O_network])
print(nn)
###############################################################################
# Initialize the weights and biases.
#
# .. note::
# Pytorch default initialization for the weights and biases in linear layers
# is Kaiming uniform. See: `TORCH.NN.MODULES.LINEAR`_
# We initialize the weights similarly but from the normal distribution.
# The biases were initialized to zero.
#
# .. _TORCH.NN.MODULES.LINEAR:
# https://pytorch.org/docs/stable/_modules/torch/nn/modules/linear.html#Linear
device = torch.device(parser.device)
builtins = torchani.neurochem.Builtins()
consts = builtins.consts
aev_computer = builtins.aev_computer
shift_energy = builtins.energy_shifter
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
model = torchani.ANIModel([atomic() for _ in range(4)])
class Flatten(torch.nn.Module):
def forward(self, x):
return x[0], x[1].flatten()
nnp = torch.nn.Sequential(model, Flatten()).to(device)
dataset = torchani.data.AEVCacheLoader(parser.cache_path)
container = torchani.ignite.Container({'energies': nnp})
optimizer = torch.optim.Adam(nnp.parameters())
trainer = ignite.engine.create_supervised_trainer(
container, optimizer, torchani.ignite.MSELoss('energies'))
def benchmark(parser, dataset, use_cuda_extension, force_inference=False):
synchronize = True
timers = {}
def time_func(key, func):
timers[key] = 0
def wrapper(*args, **kwargs):
start = timeit.default_timer()
ret = func(*args, **kwargs)
sync_cuda(synchronize)
end = timeit.default_timer()
timers[key] += end - start
return ret
return wrapper
Rcr = 5.2000e+00
Rca = 3.5000e+00
EtaR = torch.tensor([1.6000000e+01], device=parser.device)
ShfR = torch.tensor([
9.0000000e-01, 1.1687500e+00, 1.4375000e+00, 1.7062500e+00,
1.9750000e+00, 2.2437500e+00, 2.5125000e+00, 2.7812500e+00,
3.0500000e+00, 3.3187500e+00, 3.5875000e+00, 3.8562500e+00,
4.1250000e+00, 4.3937500e+00, 4.6625000e+00, 4.9312500e+00
],
device=parser.device)
Zeta = torch.tensor([3.2000000e+01], device=parser.device)
ShfZ = torch.tensor([
1.9634954e-01, 5.8904862e-01, 9.8174770e-01, 1.3744468e+00,
1.7671459e+00, 2.1598449e+00, 2.5525440e+00, 2.9452431e+00
],
device=parser.device)
EtaA = torch.tensor([8.0000000e+00], device=parser.device)
ShfA = torch.tensor(
[9.0000000e-01, 1.5500000e+00, 2.2000000e+00, 2.8500000e+00],
device=parser.device)
num_species = 4
aev_computer = torchani.AEVComputer(Rcr, Rca, EtaR, ShfR, EtaA, Zeta, ShfA,
ShfZ, num_species, use_cuda_extension)
nn = torchani.ANIModel(build_network())
model = torch.nn.Sequential(aev_computer, nn).to(parser.device)
optimizer = torch.optim.Adam(model.parameters(), lr=0.000001)
mse = torch.nn.MSELoss(reduction='none')
# enable timers
torchani.aev.cutoff_cosine = time_func('torchani.aev.cutoff_cosine',
torchani.aev.cutoff_cosine)
torchani.aev.radial_terms = time_func('torchani.aev.radial_terms',
torchani.aev.radial_terms)
torchani.aev.angular_terms = time_func('torchani.aev.angular_terms',
torchani.aev.angular_terms)
torchani.aev.compute_shifts = time_func('torchani.aev.compute_shifts',
torchani.aev.compute_shifts)
torchani.aev.neighbor_pairs = time_func('torchani.aev.neighbor_pairs',
torchani.aev.neighbor_pairs)
torchani.aev.neighbor_pairs_nopbc = time_func(
'torchani.aev.neighbor_pairs_nopbc', torchani.aev.neighbor_pairs_nopbc)
torchani.aev.triu_index = time_func('torchani.aev.triu_index',
torchani.aev.triu_index)
torchani.aev.cumsum_from_zero = time_func('torchani.aev.cumsum_from_zero',
torchani.aev.cumsum_from_zero)
torchani.aev.triple_by_molecule = time_func(
'torchani.aev.triple_by_molecule', torchani.aev.triple_by_molecule)
torchani.aev.compute_aev = time_func('torchani.aev.compute_aev',
torchani.aev.compute_aev)
model[0].forward = time_func('total', model[0].forward)
model[1].forward = time_func('forward', model[1].forward)
optimizer.step = time_func('optimizer.step', optimizer.step)
print('=> start training')
start = time.time()
loss_time = 0
force_time = 0
for epoch in range(0, parser.num_epochs):
print('Epoch: %d/%d' % (epoch + 1, parser.num_epochs))
progbar = pkbar.Kbar(target=len(dataset) - 1, width=8)
for i, properties in enumerate(dataset):
species = properties['species'].to(parser.device)
coordinates = properties['coordinates'].to(
parser.device).float().requires_grad_(force_inference)
true_energies = properties['energies'].to(parser.device).float()
num_atoms = (species >= 0).sum(dim=1, dtype=true_energies.dtype)
_, predicted_energies = model((species, coordinates))
# TODO add sync after aev is done
sync_cuda(synchronize)
energy_loss = (mse(predicted_energies, true_energies) /
num_atoms.sqrt()).mean()
if force_inference:
sync_cuda(synchronize)
force_coefficient = 0.1
true_forces = properties['forces'].to(parser.device).float()
force_start = time.time()
try:
sync_cuda(synchronize)
forces = -torch.autograd.grad(predicted_energies.sum(),
coordinates,
create_graph=True,
retain_graph=True)[0]
sync_cuda(synchronize)
except Exception as e:
alert('Error: {}'.format(e))
return
force_time += time.time() - force_start
force_loss = (mse(true_forces, forces).sum(dim=(1, 2)) /
num_atoms).mean()
loss = energy_loss + force_coefficient * force_loss
sync_cuda(synchronize)
else:
loss = energy_loss
rmse = hartree2kcalmol(
(mse(predicted_energies,
true_energies)).mean()).detach().cpu().numpy()
progbar.update(i, values=[("rmse", rmse)])
if not force_inference:
sync_cuda(synchronize)
loss_start = time.time()
loss.backward()
# print('2', coordinates.grad)
sync_cuda(synchronize)
loss_stop = time.time()
loss_time += loss_stop - loss_start
optimizer.step()
sync_cuda(synchronize)
checkgpu()
sync_cuda(synchronize)
stop = time.time()
print('=> More detail about benchmark PER EPOCH')
total_time = (stop - start) / parser.num_epochs
loss_time = loss_time / parser.num_epochs
force_time = force_time / parser.num_epochs
opti_time = timers['optimizer.step'] / parser.num_epochs
forward_time = timers['forward'] / parser.num_epochs
aev_time = timers['total'] / parser.num_epochs
print_timer(' Total AEV', aev_time)
print_timer(' Forward', forward_time)
print_timer(' Backward', loss_time)
print_timer(' Force', force_time)
print_timer(' Optimizer', opti_time)
print_timer(
' Others', total_time - loss_time - aev_time - forward_time -
opti_time - force_time)
print_timer(' Epoch time', total_time)
