def train(model="sch", epochs=80, device=th.device("cpu")): print("start") alchemy_dataset = TencentAlchemyDataset() alchemy_loader = DataLoader(dataset=alchemy_dataset, batch_size=20, collate_fn=batcher(), shuffle=False, num_workers=0) if model == "sch": model = SchNetModel(norm=True, output_dim=12) elif model == "mgcn": model = MGCNModel(norm=True, output_dim=12) elif model == "MPNN": model = MPNNModel(output_dim=12) print(model) if model.name in ["MGCN", "SchNet"]: model.set_mean_std(alchemy_dataset.mean, alchemy_dataset.std, device) model.to(device) loss_fn = nn.MSELoss() MAE_fn = nn.L1Loss() optimizer = th.optim.Adam(model.parameters(), lr=0.0001) for epoch in range(epochs): w_loss, w_mae = 0, 0 model.train() for idx, batch in enumerate(alchemy_loader): batch.graph.to(device) batch.label = batch.label.to(device) res = model(batch.graph) loss = loss_fn(res, batch.label) mae = MAE_fn(res, batch.label) optimizer.zero_grad() loss.backward() optimizer.step() w_mae += mae.detach().item() w_loss += loss.detach().item() w_mae /= idx + 1 print("Epoch {:2d}, loss: {:.7f}, mae: {:.7f}".format( epoch, w_loss, w_mae))
def sch_train(epochs=80, device=torch.device('cpu')): alchemy_dataset = TencentAlchemyDataset() alchemy_loader = DataLoader(dataset=alchemy_dataset, batch_size=20, collate_fn=batcher(), shuffle=False, num_workers=0) model = SchNetModel(norm=True, output_dim=15) model.set_mean_std(alchemy_dataset.mean, alchemy_dataset.std, device) model.to(device) loss_fn = nn.MSELoss() MAE_fn = nn.L1Loss() optimizer = torch.optim.Adam(model.parameters(), lr=0.0001) losses = [] for epoch in range(epochs): start = time() w_loss, w_mae = 0, 0 model.train() for idx, batch in enumerate(alchemy_loader): batch.graph.to(device) batch.label = batch.label.to(device) res = model(batch.graph) loss = loss_fn(res, batch.label) mae = MAE_fn(res, batch.label) optimizer.zero_grad() loss.backward() optimizer.step() w_mae += mae.detach().item() w_loss += loss.detach().item() w_mae /= idx + 1 print( "Epoch {:2d}, \tloss: {:.7f}, \tmae: {:.7f}, \ttime: {:.3f}".format( epoch, w_loss, w_mae, time() - start)) losses.append(w_loss) torch.save(model, r'./model/schnet-{}.pth'.format(epochs)) plt.plot(range(epochs), losses, color='red') plt.savefig( r'./img/schnet_{}_{:.4f}.png'.format(strftime('%Y-%m-%d'), str(w_mae))) plt.show()
def eval(model="sch", epochs=80, device=th.device("cpu"), train_dataset='', eval_dataset='', epoch=1): print("start") epoch = int(epoch) test_dataset = TencentAlchemyDataset() test_dir = './' test_file = train_dataset + '_' + eval_dataset + "_cross.csv" test_dataset.mode = "Train" test_dataset.transform = None test_dataset.file_path = test_file test_dataset._load() test_loader = DataLoader( dataset=test_dataset, batch_size=10, collate_fn=batcher(), shuffle=False, num_workers=0, ) if model == "sch": model = SchNetModel(norm=False, output_dim=1) elif model == "mgcn": model = MGCNModel(norm=False, output_dim=1) elif model == "MPNN": model = MPNNModel(output_dim=1) print(model) # if model.name in ["MGCN", "SchNet"]: # model.set_mean_std(mean, std, device) model.load_state_dict( th.load('./' + train_dataset + "/model_" + str(epoch))) model.to(device) loss_fn = nn.MSELoss() MAE_fn = nn.L1Loss() # optimizer = th.optim.Adam(model.parameters(), lr=0.0001) val_loss, val_mae = 0, 0 res_file = open( train_dataset + '_' + eval_dataset + str(epoch) + "_crossres.txt", 'w') for jdx, batch in enumerate(test_loader): batch.graph.to(device) batch.label = batch.label.to(device) res = model(batch.graph) res_np = res.cpu().detach().numpy() label_np = batch.label.cpu().detach().numpy() for i in range(len(res_np)): res_file.write(str(res_np[i][0]) + '\t') res_file.write(str(label_np[i][0]) + '\n') loss = loss_fn(res, batch.label) mae = MAE_fn(res, batch.label) # optimizer.zero_grad() # loss.backward() # optimizer.step() val_mae += mae.detach().item() val_loss += loss.detach().item() val_mae /= jdx + 1 val_loss /= jdx + 1 print("Epoch {:2d}, val_loss: {:.7f}, val_mae: {:.7f}".format( epoch, val_loss, val_mae)) print("test_dataset.mean= %s" % (test_dataset.mean)) print("test_dataset.std= %s" % (test_dataset.std))
def train(model="sch", epochs=80, device=th.device("cpu"), dataset='', save=''): print("start") train_dir = "./" train_file = dataset + "_train.csv" alchemy_dataset = TencentAlchemyDataset() alchemy_dataset.mode = "Train" alchemy_dataset.transform = None alchemy_dataset.file_path = train_file alchemy_dataset._load() test_dataset = TencentAlchemyDataset() test_dir = train_dir test_file = dataset + "_valid.csv" test_dataset.mode = "Train" test_dataset.transform = None test_dataset.file_path = test_file test_dataset._load() alchemy_loader = DataLoader( dataset=alchemy_dataset, batch_size=10, collate_fn=batcher(), shuffle=False, num_workers=0, ) test_loader = DataLoader( dataset=test_dataset, batch_size=10, collate_fn=batcher(), shuffle=False, num_workers=0, ) if model == "sch": model = SchNetModel(norm=False, output_dim=1) elif model == "mgcn": model = MGCNModel(norm=False, output_dim=1) print(model) # if model.name in ["MGCN", "SchNet"]: # model.set_mean_std(alchemy_dataset.mean, alchemy_dataset.std, device) model.to(device) # print("test_dataset.mean= %s" % (alchemy_dataset.mean)) # print("test_dataset.std= %s" % (alchemy_dataset.std)) loss_fn = nn.MSELoss() MAE_fn = nn.L1Loss() optimizer = th.optim.Adam(model.parameters(), lr=0.0001) scheduler = th.optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', factor=0.9, patience=10, threshold=0.0000001, threshold_mode='rel', cooldown=0, min_lr=0.000001, eps=1e-08, verbose=False) for epoch in range(epochs): w_loss, w_mae = 0, 0 model.train() for idx, batch in enumerate(alchemy_loader): batch.graph.to(device) batch.label = batch.label.to(device) res = model(batch.graph) loss = loss_fn(res, batch.label) mae = MAE_fn(res, batch.label) optimizer.zero_grad() loss.backward() optimizer.step() w_mae += mae.detach().item() w_loss += loss.detach().item() w_mae /= idx + 1 w_loss /= idx + 1 scheduler.step(w_mae) print("Epoch {:2d}, loss: {:.7f}, mae: {:.7f}".format( epoch, w_loss, w_mae)) val_loss, val_mae = 0, 0 for jdx, batch in enumerate(test_loader): batch.graph.to(device) batch.label = batch.label.to(device) res = model(batch.graph) loss = loss_fn(res, batch.label) mae = MAE_fn(res, batch.label) # optimizer.zero_grad() # mae.backward() # optimizer.step() val_mae += mae.detach().item() val_loss += loss.detach().item() val_mae /= jdx + 1 val_loss /= jdx + 1 print("Epoch {:2d}, val_loss: {:.7f}, val_mae: {:.7f}".format( epoch, val_loss, val_mae)) if epoch % 200 == 0: th.save(model.state_dict(), save + "/model_" + str(epoch))
def train(model="sch", epochs=80, device=th.device("cpu"), dataset=''): print("start") train_dir = "./" train_file = "train_smi.csv" alchemy_dataset = TencentAlchemyDataset() alchemy_dataset.mode = "Train" alchemy_dataset.transform = None alchemy_dataset.file_path = train_file alchemy_dataset._load() test_dataset = TencentAlchemyDataset() test_dir = train_dir test_file = "val_smi.csv" test_dataset.mode = "Train" test_dataset.transform = None test_dataset.file_path = test_file test_dataset._load() alchemy_loader = DataLoader( dataset=alchemy_dataset, batch_size=10, collate_fn=batcher(), shuffle=False, num_workers=0, ) test_loader = DataLoader( dataset=test_dataset, batch_size=10, collate_fn=batcher(), shuffle=False, num_workers=0, ) if model == "sch": model = SchNetModel(norm=False, output_dim=1) elif model == "mgcn": model = MGCNModel(norm=False, output_dim=1) elif model == "MPNN": model = MPNNModel(output_dim=1) print(model) # if model.name in ["MGCN", "SchNet"]: # model.set_mean_std(alchemy_dataset.mean, alchemy_dataset.std, device) model.to(device) # print("test_dataset.mean= %s" % (alchemy_dataset.mean)) # print("test_dataset.std= %s" % (alchemy_dataset.std)) loss_fn = nn.MSELoss() MAE_fn = nn.L1Loss() optimizer = th.optim.Adam(model.parameters(), lr=0.0001) for epoch in range(epochs): w_loss, w_mae = 0, 0 model.train() for idx, batch in enumerate(alchemy_loader): batch.graph.to(device) batch.label = batch.label.to(device) res = model(batch.graph) loss = loss_fn(res, batch.label) mae = MAE_fn(res, batch.label) optimizer.zero_grad() loss.backward() optimizer.step() w_mae += mae.detach().item() w_loss += loss.detach().item() w_mae /= idx + 1 w_loss /= idx + 1 print("Epoch {:2d}, loss: {:.7f}, mae: {:.7f}".format( epoch, w_loss, w_mae)) val_loss, val_mae = 0, 0 if (epoch % 50 == 0): res_file = open('val_results_%s.txt' % (epoch), 'w') for jdx, batch in enumerate(test_loader): batch.graph.to(device) batch.label = batch.label.to(device) res = model(batch.graph) loss = loss_fn(res, batch.label) mae = MAE_fn(res, batch.label) optimizer.zero_grad() mae.backward() optimizer.step() val_mae += mae.detach().item() val_loss += loss.detach().item() res_np = res.cpu().detach().numpy() label_np = batch.label.cpu().detach().numpy() if (epoch % 50 == 0): for i in range(len(res_np)): res_file.write(str(res_np[i][0]) + '\t') res_file.write(str(label_np[i][0]) + '\n') val_mae /= jdx + 1 val_loss /= jdx + 1 print("Epoch {:2d}, val_loss: {:.7f}, val_mae: {:.7f}".format( epoch, val_loss, val_mae)) if epoch % 50 == 0: th.save(model.state_dict(), './' + dataset + "/model_" + str(epoch))