def testSAE(self): shifter = torchani.EnergyShifter(None) torchani.data.load(dataset_path).subtract_self_energies(shifter) true_self_energies = torch.tensor([ -19.354171758844188, -19.354171758844046, -54.712238523648587, -75.162829556770987 ], dtype=torch.float64) self.assertEqual(true_self_energies, shifter.self_energies)
def testShapeInference(self): shifter = torchani.EnergyShifter(None) ds = torchani.data.load(dataset_path).subtract_self_energies(shifter) len(ds) ds = ds.species_to_indices() len(ds) ds = ds.shuffle() len(ds) ds = ds.collate(batch_size) len(ds)
def testDataloader(self): shifter = torchani.EnergyShifter(None) dataset = list( torchani.data.load(dataset_path).subtract_self_energies( shifter).species_to_indices().shuffle()) loader = torch.utils.data.DataLoader( dataset, batch_size=batch_size, collate_fn=torchani.data.collate_fn, num_workers=64) for _ in loader: pass
def testSAE(self): tolerance = 1e-5 shifter = torchani.EnergyShifter(None) torchani.data.load(dataset_path).subtract_self_energies(shifter) true_self_energies = torch.tensor([ -19.354171758844188, -19.354171758844046, -54.712238523648587, -75.162829556770987 ], dtype=torch.float64) diff = torch.abs(true_self_energies - shifter.self_energies) for e in diff: self.assertLess(e, tolerance)
def __init__(self, label="ase", atoms=None, coordinates='tmp_ase.pdb', **kwargs): self.label = label coordinates = os.path.dirname(label) + "/" + coordinates self.coordinates = coordinates self.atoms = atoms self.energy_free = None self.forces = [] self.aev_computer = torchani.SortedAEV(const_file=const_file, device=device) self.nn = torchani.ModelOnAEV(self.aev_computer, derivative=True, from_nc=network_dir, ensemble=8) self.shift_energy = torchani.EnergyShifter(sae_file)
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) nn = torchani.ANIModel([atomic() for _ in range(4)]) 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') print('=> loading dataset...') shifter = torchani.EnergyShifter(None) dataset = list( torchani.data.load(parser.dataset_path).subtract_self_energies( shifter).species_to_indices().shuffle().collate(parser.batch_size)) print('=> start warming up') total_batch_counter = 0 for epoch in range(0, WARM_UP_BATCHES + 1): print('Epoch: %d/inf' % (epoch + 1, )) progbar = pkbar.Kbar(target=len(dataset) - 1, width=8) for i, properties in enumerate(dataset): if not parser.dry_run and total_batch_counter == WARM_UP_BATCHES: print('=> warm up finished, start profiling')