mlatom 2.0.1__py3.7.egg

MLatom/interfaces/PhysNet/interface_PhysNet.py

@@ -0,0 +1,270 @@
+#!/usr/bin/python3
+'''
+  !---------------------------------------------------------------------------! 
+  ! Interface_PhysNet: Interface between PhysNet and MLatom                   ! 
+  ! Implementations by: Fuchun Ge and Max Pinheiro Jr                         ! 
+  !---------------------------------------------------------------------------! 
+'''
+import numpy as np
+import os, sys, subprocess, time, shutil, re, math, random
+import stopper
+from args_class import ArgsBase
+
+filedir = os.path.dirname(__file__)
+
+class Args(ArgsBase):
+    def __init__(self):
+        super().__init__()
+        self.add_default_dict_args([
+            'learningCurve'
+            ],
+            bool
+        )
+        self.add_default_dict_args([
+            'xyzfile', 'yfile', 'ygradxyzfile','itrainin','itestin','isubtrainin','ivalidatein','mlmodelin','setname'
+            ],
+            ""
+        )        
+        self.add_dict_args({
+            'mlmodeltype': 'PhysNet',
+            'mlmodelout': "PhysNet",
+            'sampling': "random",
+            'yestfile': "enest.dat",
+            'ygradxyzestfile': "gradest.dat",
+            'lcNtrains': [],
+            'natom': 0,
+            'atype': []
+        })
+        self.parse_input_content([
+            'physnet.batch_sizes=0',
+            'physnet.restart=0',
+            'physnet.num_features=128',
+            'physnet.num_basis=64',
+            'physnet.num_blocks=5',
+            'physnet.num_residual_atomic=2',
+            'physnet.num_residual_interaction=3',
+            'physnet.num_residual_output=1',
+            'physnet.cutoff=10.0',
+            'physnet.use_electrostatic=0',
+            'physnet.use_dispersion=0',
+            'physnet.grimme_s6=0.5',
+            'physnet.grimme_s8=0.2130',
+            'physnet.grimme_a1=0.0',
+            'physnet.grimme_a2=6.0519',
+            'physnet.num_train=8',
+            'physnet.num_valid=2',
+            'physnet.seed=42',
+            'physnet.max_steps=10000000',
+            'physnet.learning_rate=0.0008',
+            'physnet.max_norm=1000.0',
+            'physnet.ema_decay=0.999',
+            'physnet.keep_prob=1.0',
+            'physnet.l2lambda=0.0',
+            'physnet.nhlambda=0.01',
+            'physnet.decay_steps=10000000',
+            'physnet.decay_rate=0.1',
+            'physnet.batch_size=12',
+            'physnet.valid_batch_size=2',
+            'physnet.force_weight=52.91772105638412',
+            'physnet.charge_weight=0',
+            'physnet.dipole_weight=0',
+            'physnet.summary_interval=5',
+            'physnet.validation_interval=5',
+            'physnet.save_interval=10',
+            'physnet.record_run_metadata=0'
+            ])
+    def parse(self, argsraw):
+        self.parse_input_content(argsraw)
+        self.argProcess()
+
+    def argProcess(self):
+        with open(self.xyzfile,'r') as f:
+            self.natom = int(f.readline())
+            exec('f.readline()')
+            self.atype = [f.readline().split()[0] for i in range(self.natom)]
+        if not self.mlmodelin:
+            self.mlmodelin = self.mlmodelout
+        if not self.ygradxyzfile:
+            self.physnet.force_weight=0
+        
+        if self.learningcurve:
+            self.lcNtrains = [int(i) for i in str(self.lcNtrains).split(',')]
+            if self.physnet.batch_sizes:
+                self.physnet.batch_sizes = [int(i) for i in self.lcNtrains.split(',')]
+                self.physnet.batch_size = self.physnet.batch_sizes[self.lcNtrains.index(self.ntrain)]
+
+class PhysNetCls(object):
+    dataConverted = False
+
+    @classmethod
+    def load(cls):
+        loaded=False
+        if not loaded:
+            
+            sp_z = {'X': 0, 'H': 1, 'He': 2, 'Li': 3, 'Be': 4, 'B': 5, 'C': 6, 'N': 7, 'O': 8, 'F': 9, 'Ne': 10, 'Na': 11, 'Mg': 12, 'Al': 13, 'Si': 14, 'P': 15, 'S': 16, 'Cl': 17, 'Ar': 18, 'K': 19, 'Ca': 20, 'Sc': 21, 'Ti': 22, 'V': 23, 'Cr': 24, 'Mn': 25, 'Fe': 26, 'Co': 27, 'Ni': 28, 'Cu': 29, 'Zn': 30, 'Ga': 31, 'Ge': 32, 'As': 33, 'Se': 34, 'Br': 35, 'Kr': 36, 'Rb': 37, 'Sr': 38, 'Y': 39, 'Zr': 40, 'Nb': 41, 'Mo': 42, 'Tc': 43, 'Ru': 44, 'Rh': 45, 'Pd': 46, 'Ag': 47, 'Cd': 48, 'In': 49, 'Sn': 50, 'Sb': 51, 'Te': 52, 'I': 53, 'Xe': 54, 'Cs': 55, 'Ba': 56, 'La': 57, 'Ce': 58, 'Pr': 59, 'Nd': 60, 'Pm': 61, 'Sm': 62, 'Eu': 63, 'Gd': 64, 'Tb': 65, 'Dy': 66, 'Ho': 67, 'Er': 68, 'Tm': 69, 'Yb': 70, 'Lu': 71, 'Hf': 72, 'Ta': 73, 'W': 74, 'Re': 75, 'Os': 76, 'Ir': 77, 'Pt': 78, 'Au': 79, 'Hg': 80, 'Tl': 81, 'Pb': 82, 'Bi': 83, 'Po': 84, 'At': 85, 'Rn': 86, 'Fr': 87, 'Ra': 88, 'Ac': 89, 'Th': 90, 'Pa': 91, 'U': 92, 'Np': 93, 'Pu': 94, 'Am': 95, 'Cm': 96, 'Bk': 97, 'Cf': 98, 'Es': 99, 'Fm': 100, 'Md': 101, 'No': 102, 'Lr': 103, 'Rf': 104, 'Db': 105, 'Sg': 106, 'Bh': 107, 'Hs': 108, 'Mt': 109, 'Ds': 110, 'Rg': 111, 'Cn': 112, 'Uut': 113, 'Fl': 114, 'Uup': 115, 'Lv': 116, 'Uus': 117, 'Uuo': 118}
+
+            z_sp = {v:k for k,v in sp_z.items()}
+            try:
+                from . import PhysNet_train
+                from . import PhysNet_predict
+            except:
+                stopper.stopMLatom('Please specify PhysNet installation dir in $PhysNet')
+
+            globals()['z_sp'] = z_sp
+            globals()['sp_z'] = sp_z
+            globals()['PhysNet_train'] = PhysNet_train
+            globals()['PhysNet_predict'] = PhysNet_predict
+            loaded=True
+
+
+    def __init__(self, argsPhysNet = sys.argv[1:]):
+        print(' ___________________________________________________________\n\n%s' % __doc__)
+    
+    @classmethod
+    def convertdata(cls, argsPhysNet, subdatasets):
+        cls.load()
+        def convert(fileout, setname, yorgrad=False):
+            # convert data to .h5 format
+            prefix = ''
+            if args.learningcurve: prefix = '../'
+            if setname: 
+                coordfile = prefix+'xyz.dat_'+setname
+                yfile = prefix+'y.dat_'+setname
+                gradfile = prefix+'grad.dat_'+setname
+            else: 
+                coordfile = args.xyzfile
+                if args.yfile and yorgrad: yfile = args.yfile
+                if args.ygradxyzfile and yorgrad: gradfile = args.ygradxyzfile
+            with open(coordfile ,'r') as fxyz:
+                natom = int(fxyz.readline())
+                fxyz.seek(0)
+                data={}
+                data['R'] = np.array([])
+                data['Z'] = np.array([])
+                data['N'] = np.array([])
+                if args.yfile and yorgrad: 
+                    fy =  open(yfile,'r')
+                    data['E'] = np.array([])
+                if args.ygradxyzfile and yorgrad: 
+                    fgrad = open(gradfile,'r')
+                    data['F'] = np.array([])
+                for line in fxyz:
+                    # 1st line
+                    data['N'] = np.append(data['N'], int(line)).astype('int')
+                    if args.ygradxyzfile and yorgrad: fgrad.readline()
+                    if args.yfile and yorgrad: data['E'] =np.append(data['E'], float(fy.readline()))
+                    # 2nd line
+                    fxyz.readline()
+                    if args.ygradxyzfile and yorgrad: fgrad.readline()
+                    # rest lines
+                    for i in range(natom):
+                        ln = fxyz.readline()
+                        data['Z'] = np.append(data['Z'], sp_z[ln.split()[0].capitalize()])
+                        data['R'] = np.append(data['R'], [float(i) for i in ln.split()[-3:]])
+                        if args.ygradxyzfile and yorgrad: data['F'] = np.append(data['F'], [-1*float(i) for i in fgrad.readline().split()[-3:]]) 
+
+                if args.yfile and yorgrad: fy.close()
+                if args.ygradxyzfile and yorgrad: fgrad.close()
+                data['R'] = data['R'].reshape(-1,natom,3)
+                data['Z'] = data['Z'].reshape(-1,natom).astype('int')
+                if args.ygradxyzfile and yorgrad: data['F'] = data['F'].reshape(-1,natom,3)
+                
+                np.savez(fileout, **data)
+
+            
+        args = Args()
+        args.parse(argsPhysNet)
+        print('\n Converting data...\n\n')
+        if subdatasets:
+            if 'Subtrain'  in subdatasets:
+                for i in subdatasets:
+                    if i in ['Subtrain', 'Validate']: convert('PhysNet_'+i.lower()+'.npz', i.lower(),True)
+                    else: convert('PhysNet_'+i.lower()+'.npz', i.lower())
+            else:
+                for i in subdatasets:
+                    if i == 'Train': convert('PhysNet_'+i.lower()+'.npz', i.lower(),True)
+                    else: convert('PhysNet_'+i.lower()+'.npz', i.lower())
+        else:   convert('PhysNet.npz', '')
+
+    @classmethod
+    def createMLmodel(cls, argsPhysNet, subdatasets):
+        cls.load()
+        PNargs = Args()
+        PNargs.parse(argsPhysNet)
+
+        if not cls.dataConverted or PNargs.learningcurve:
+            cls.convertdata(argsPhysNet, subdatasets)
+            cls.dataConverted = True
+        
+        starttime = time.time()
+
+        PhysNet_train.train(PNargs)
+
+        endtime = time.time()
+        wallclock = endtime - starttime
+        return wallclock
+        
+    @classmethod
+    def useMLmodel(cls, argsPhysNet, subdatasets):
+        cls.load()
+        PNargs = Args()
+        PNargs.parse(argsPhysNet)
+
+        if not cls.dataConverted: 
+            cls.convertdata(argsPhysNet, subdatasets)
+            cls.dataConverted = True
+
+        starttime = time.time()
+        PhysNet_predict.predict(PNargs)
+        endtime = time.time()
+        wallclock = endtime - starttime
+        return wallclock
+
+def printHelp():
+    helpText = __doc__ + '''
+  To use Interface_PhysNet, please define $PhysNet to where PhysNet is located
+
+  Arguments with their default values:
+    MLprog=PhysNet             enables this interface
+    MLmodelType=PhysNet        requests PhysNet model
+    
+    physnet.num_features=128   number of input features
+    physnet.num_basis=64       number of radial basis functions
+    physnet.num_blocks=5       number of stacked modular building blocks
+                                                
+    physnet.num_residual_atomic=2 
+                               number of residual blocks for 
+                               atom-wise refinements
+    physnet.num_residual_interaction=3
+                               number of residual blocks for 
+                               refinements of proto-message
+    physnet.num_residual_output=1 
+                               number of residual blocks in 
+                               output blocks
+    physnet.cutoff=10.0        cutoff radius for interactions 
+                               in the neural network
+    physnet.seed=42            random seed
+    physnet.max_steps=32       max steps to perform in training
+    physnet.learning_rate=0.0008 
+                               starting learning rate
+    physnet.decay_steps=10000000 
+                               decay steps
+    physnet.decay_rate=0.1     decay rate for learning rate
+    physnet.batch_size=12      training batch size
+    physnet.valid_batch_size=2 validation batch size
+    physnet.force_weight=52.91772105638412 
+                               weight for force
+    physnet.charge_weight=0    weight for charge
+    physnet.dipole_weight=0    weight for dipole
+    physnet.summary_interval=5 interval for summary
+    physnet.validation_interval=5 
+                               interval for validation
+    physnet.save_interval=10   interval for model saving
+
+  Cite PhysNet:
+    O. T. Unke, M. Meuwly, J. Chem. Theory Comput. 2019, 15, 3678
+'''
+    print(helpText)
+
+if __name__ == '__main__':
+    PhysNetCls()

MLatom/interfaces/TorchANI/TorchANI_predict.py

@@ -0,0 +1,82 @@
+import torch
+import torchani
+import os
+import math
+import torch.utils.tensorboard
+import tqdm
+import numpy as np
+import h5py
+
+def predict(args):
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+
+    species_order = sorted(set(args.atype), key=lambda x: args.atype.index(x))
+    energyOffset = 0
+
+    # if args.mlmodeltype.lower() == 'ani1x':
+    #     model = torchani.models.ANI1x(periodic_table_index=True).to(device)
+    # elif args.mlmodeltype.lower() == 'ani1ccx':
+    #     model = torchani.models.ANI1ccx(periodic_table_index=True).to(device)
+    # elif args.mlmodeltype.lower() == 'ani2x':
+    #     model = torchani.models.ANI2x(periodic_table_index=True).to(device)
+    if False:
+        pass
+    else:
+        best_model = torch.load(args.mlmodelin)
+
+        Rcr = best_model['args']['Rcr']
+        Rca = best_model['args']['Rca']
+        EtaR = best_model['args']['EtaR'].to(device)
+        ShfR = best_model['args']['ShfR'].to(device)
+        Zeta = best_model['args']['Zeta'].to(device)
+        ShfZ = best_model['args']['ShfZ'].to(device)
+        EtaA = best_model['args']['EtaA'].to(device)
+        ShfA = best_model['args']['ShfA'].to(device)
+
+
+        num_species = len(species_order)
+
+        self_energies=best_model['args']['self_energies']
+        energyOffset=best_model['args']['energyOffset']
+
+        aev_computer = torchani.AEVComputer(Rcr, Rca, EtaR, ShfR, EtaA, Zeta, ShfA, ShfZ, num_species)
+        aev_dim = aev_computer.aev_length
+        networkdic = best_model['network']
+
+        nn = torchani.ANIModel([networkdic[specie] for specie in species_order])
+        nn.load_state_dict(best_model['nn'])
+        
+        model = torchani.nn.Sequential(aev_computer, nn).to(device)
+    model.eval()
+
+    batch_size = args.ani.batch_size
+                
+    if args.setname: args.setname='_'+args.setname
+    testfile=h5py.File('ANI'+args.setname+'.h5','r')
+    test=testfile.get('dataset').get('molecule')
+    all_coordinates=torch.tensor(test.get('coordinates')[()]).to(device).float()
+
+    with open(args.yestfile,'wb') as fy,  open(args.ygradxyzestfile,'wb') as fgrad:
+        for i in range(math.ceil(len(all_coordinates)/batch_size)):
+            coordinates=all_coordinates[batch_size*i:min(batch_size*(i+1),len(all_coordinates))].requires_grad_(True)
+
+            species=test.get('species')[()]
+            species=torch.tensor([[species_order.index(i.decode('ascii')) for i in species] for j in range(len(coordinates))]).to(device)
+
+            _, predicted_energies = model((species, coordinates))
+            predicted_forces = -torch.autograd.grad(predicted_energies.sum(), coordinates, create_graph=True, retain_graph=True)[0]
+            if args.ygradxyzfile: 
+                predicted_energies = np.array(predicted_energies.cpu().detach().numpy()).reshape(-1,1)+energyOffset
+                predicted_forces = np.array(predicted_forces.cpu().detach().numpy()).reshape(-1, args.natom, 3)
+            else: 
+                predicted_energies = np.array(predicted_energies.cpu().detach().numpy()).reshape(-1,1)+energyOffset
+                predicted_forces = np.array(predicted_forces.cpu().detach().numpy()).reshape(-1, args.natom, 3)
+                        
+            np.savetxt(fy, predicted_energies, fmt='%20.12f', delimiter=" ")
+            line="%d\n\n" % args.natom
+            for force in predicted_forces:
+                fgrad.write(line.encode('utf-8'))
+                np.savetxt(fgrad, -1*force, fmt='%20.12f', delimiter=" ")
+                
+    import time
+    time.sleep(0.0000001) # Sometimes program hangs without no reason, adding short sleep time helps to exit this module without a problem / P.O.D., 2021-02-17

MLatom/interfaces/TorchANI/TorchANI_train.py

@@ -0,0 +1,217 @@
+import torch
+import torchani
+import os
+import sys
+import math
+import torch.utils.tensorboard
+import tqdm
+import numpy as np
+
+def train(args):
+
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    print('Training with %s'% device)
+
+    Rcr = args.ani.Rcr
+    Rca = args.ani.Rca
+    EtaR = args.ani.EtaR.to(device)
+    ShfR = args.ani.ShfR.to(device)
+    Zeta = args.ani.Zeta.to(device)
+    ShfZ = args.ani.ShfZ.to(device)
+    EtaA = args.ani.EtaA.to(device)
+    ShfA = args.ani.ShfA.to(device)
+
+
+
+    max_epochs = args.ani.max_epochs
+    early_stopping_learning_rate = args.ani.early_stopping_learning_rate
+    if args.ygradxyzfile:
+        force_coefficient = args.ani.force_coefficient
+    best_model_checkpoint = args.mlmodelout 
+
+
+    species_order = sorted(set(args.atype), key=lambda x: args.atype.index(x))
+    num_species = len(species_order)
+    energy_shifter = torchani.utils.EnergyShifter(None)
+
+
+    trainfile = 'ANI_subtrain.h5'
+    validaefile = 'ANI_validate.h5'
+
+    batch_size = args.ani.batch_size
+    if args.ygradxyzfile:
+        training = torchani.data.load(trainfile, additional_properties=('forces',)).subtract_self_energies(energy_shifter, species_order).species_to_indices(species_order).shuffle()
+        energyOffset=sum([np.array(energy_shifter.self_energies)[species_order.index(i)] for i in args.atype])
+        validation= torchani.data.load(validaefile, additional_properties=('forces',)).subtract_self_energies(energy_shifter, species_order).species_to_indices(species_order).shuffle()
+    else:
+        training = torchani.data.load(trainfile).subtract_self_energies(energy_shifter, species_order).species_to_indices(species_order).shuffle()
+        energyOffset=sum([np.array(energy_shifter.self_energies)[species_order.index(i)] for i in args.atype])
+        validation = torchani.data.load(validaefile).subtract_self_energies(energy_shifter, species_order).species_to_indices(species_order).shuffle()
+    
+    training = training.collate(batch_size).cache()
+    validation = validation.collate(batch_size).cache()
+
+    # print('Self atomic energies: ', energy_shifter.self_energies)
+    
+    energyOffset1=sum([np.array(energy_shifter.self_energies)[species_order.index(i)] for i in args.atype])
+
+    deltaOffset=energyOffset-energyOffset1
+
+    argdic = {'Rcr': Rcr, 'Rca': Rca, 'EtaR': EtaR, 'ShfR': ShfR, 'Zeta': Zeta, 'ShfZ': ShfZ, 'EtaA': EtaA, 'ShfA': ShfA, 'self_energies': np.array(energy_shifter.self_energies),'energyOffset':energyOffset}
+
+
+    aev_computer = torchani.AEVComputer(Rcr, Rca, EtaR, ShfR, EtaA, Zeta, ShfA, ShfZ, num_species)
+    aev_dim = aev_computer.aev_length
+    networkdic ={}
+    for i, specie in enumerate(species_order):
+        networkdic[specie] = torch.nn.Sequential(
+            torch.nn.Linear(aev_dim, int(args.ani.Neuron_l1[i])),
+            torch.nn.CELU(0.1),
+            torch.nn.Linear(int(args.ani.Neuron_l1[i]), int(args.ani.Neuron_l2[i])),
+            torch.nn.CELU(0.1),
+            torch.nn.Linear(int(args.ani.Neuron_l2[i]), int(args.ani.Neuron_l3[i])),
+            torch.nn.CELU(0.1),
+            torch.nn.Linear(int(args.ani.Neuron_l3[i]), 1)
+        )
+    nn = torchani.ANIModel([networkdic[specie] for specie in species_order])
+
+    def init_params(m):
+        if isinstance(m, torch.nn.Linear):
+            torch.nn.init.kaiming_normal_(m.weight, a=1.0)
+            torch.nn.init.zeros_(m.bias)
+
+    nn.apply(init_params)
+    
+    print('batch_size: %s' % batch_size)
+
+    print('Neural Network Structure:\n', nn)
+
+
+    model = torchani.nn.Sequential(aev_computer, nn).to(device)
+
+    wlist2d = [
+        [
+            {'params': [networkdic[specie][0].weight]},
+            {'params': [networkdic[specie][2].weight], 'weight_decay': 0.00001},
+            {'params': [networkdic[specie][4].weight], 'weight_decay': 0.000001},
+            {'params': [networkdic[specie][6].weight]},
+        ]
+        for specie in species_order
+    ]
+
+    AdamW = torch.optim.AdamW([i for j in wlist2d for i in j])
+
+    blist2d = [
+        [
+            {'params': [networkdic[specie][0].bias]},
+            {'params': [networkdic[specie][2].bias]},
+            {'params': [networkdic[specie][4].bias]},
+            {'params': [networkdic[specie][6].bias]},
+        ]
+        for specie in species_order
+    ]
+
+    SGD = torch.optim.SGD([i for j in blist2d for i in j], lr=1e-3)
+  
+    AdamW_scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(AdamW, factor=args.ani.lrfactor, patience=args.ani.patience, threshold=0)
+    SGD_scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(SGD, factor=args.ani.lrfactor, patience=args.ani.patience, threshold=0)    
+
+    latest_checkpoint = 'latest.pt'
+
+    if os.path.isfile(latest_checkpoint):
+        checkpoint = torch.load(latest_checkpoint)
+        nn.load_state_dict(checkpoint['nn'])
+        AdamW.load_state_dict(checkpoint['AdamW'])
+        SGD.load_state_dict(checkpoint['SGD'])
+        AdamW_scheduler.load_state_dict(checkpoint['AdamW_scheduler'])
+        SGD_scheduler.load_state_dict(checkpoint['SGD_scheduler'])
+
+    def validate():
+        # run validation
+        mse_sum = torch.nn.MSELoss(reduction='sum')
+        total_mse = 0.0
+        count = 0
+        for properties in validation:
+            species = properties['species'].to(device)
+            coordinates = properties['coordinates'].to(device).float()
+            true_energies = properties['energies'].to(device).float()
+            _, predicted_energies = model((species, coordinates))
+            # print(true_energies-deltaOffset,'\n\n',predicted_energies)
+            total_mse += mse_sum(predicted_energies, true_energies-deltaOffset).item()
+            count += predicted_energies.shape[0]
+        return math.sqrt(total_mse / count)
+
+    tensorboard = torch.utils.tensorboard.SummaryWriter()
+
+    mse = torch.nn.MSELoss(reduction='none')
+
+    print("training starting from epoch", AdamW_scheduler.last_epoch + 1)
+
+    for _ in range(AdamW_scheduler.last_epoch + 1, max_epochs):
+        rmse = validate()
+        print('RMSE:', rmse, 'at epoch', AdamW_scheduler.last_epoch + 1)
+        sys.stdout.flush()
+
+        learning_rate = AdamW.param_groups[0]['lr']
+        print('learning_rate:',learning_rate)
+
+        if learning_rate < early_stopping_learning_rate:
+            break
+
+        if AdamW_scheduler.is_better(rmse, AdamW_scheduler.best):
+            torch.save(
+                {   
+                    'network': networkdic,
+                    'args': argdic,
+                    'nn': nn.state_dict()
+                }
+                , best_model_checkpoint
+            )
+            print('best model saved')
+        AdamW_scheduler.step(rmse)
+        SGD_scheduler.step(rmse)
+
+        tensorboard.add_scalar('validation_rmse', rmse, AdamW_scheduler.last_epoch)
+        tensorboard.add_scalar('best_validation_rmse', AdamW_scheduler.best, AdamW_scheduler.last_epoch)
+        tensorboard.add_scalar('learning_rate', learning_rate, AdamW_scheduler.last_epoch)
+        for i, properties in tqdm.tqdm(
+            enumerate(training),
+            total=len(training),
+            desc="epoch {}".format(AdamW_scheduler.last_epoch)
+        ):
+            species = properties['species'].to(device)
+
+            if args.ygradxyzfile: 
+                coordinates = properties['coordinates'].to(device).float().requires_grad_(True)
+                true_energies = properties['energies'].to(device).float()
+                true_forces = properties['forces'].to(device).float()
+                num_atoms = (species >= 0).sum(dim=1, dtype=true_energies.dtype)
+                _, predicted_energies = model((species, coordinates))
+                forces = -torch.autograd.grad(predicted_energies.sum(), coordinates, create_graph=True, retain_graph=True)[0]
+                energy_loss = (mse(predicted_energies, true_energies) / num_atoms.sqrt()).mean()
+                force_loss = (mse(true_forces, forces).sum(dim=(1, 2)) / num_atoms).mean()
+                loss = energy_loss + force_coefficient * force_loss
+                # print('\n\n',true_energies+energyOffset,'\n\n',predicted_energies+energyOffset)
+            else:
+                coordinates = properties['coordinates'].to(device).float()
+                true_energies = properties['energies'].to(device).float()
+                num_atoms = (species >= 0).sum(dim=1, dtype=true_energies.dtype)
+                _, predicted_energies = model((species, coordinates))
+                loss = (mse(predicted_energies, true_energies) / num_atoms.sqrt()).mean()
+
+            AdamW.zero_grad()
+            SGD.zero_grad()
+            loss.backward()
+            AdamW.step()
+            SGD.step()
+
+            tensorboard.add_scalar('batch_loss', loss, AdamW_scheduler.last_epoch * len(training) + i)
+
+        torch.save({
+            'nn': nn.state_dict(),
+            'AdamW': AdamW.state_dict(),
+            'SGD': SGD.state_dict(),
+            'AdamW_scheduler': AdamW_scheduler.state_dict(),
+            'SGD_scheduler': SGD_scheduler.state_dict(),
+        }, latest_checkpoint)
+