hjxdl 0.0.1__py3-none-any.whl

hdl/controllers/train/rxn_train.py

@@ -0,0 +1,219 @@
+from os import path as osp
+import typing as t
+
+import torch
+from torch import nn
+
+from hdl.models.rxn import build_rxn_mu
+from hdl.models.utils import load_model, save_model
+from hdl.data.dataset.seq.rxn_dataset import RXNCSVDataset
+from hdl.data.dataset.loaders.rxn_loader import RXNLoader
+from hdl.metric_loss.loss import mtmc_loss
+from jupyfuncs.pbar import tnrange, tqdm
+from jupyfuncs.glob import makedirs
+# from hdl.optims.nadam import Nadam
+from torch.optim import Adam
+# from .trainer_base import TorchTrainer
+
+
+def train_a_batch(
+    model,
+    batch_data,
+    loss_func,
+    optimizer,
+    device,
+    individual,
+    **kwargs
+):
+    optimizer.zero_grad()
+
+    X = [x.to(device) for x in batch_data[0]]
+    y = batch_data[1].T.to(device)
+
+    y_preds = model(X)
+    loss = mtmc_loss(
+        y_preds,
+        y,
+        loss_func,
+        individual=individual, **kwargs
+    )
+
+    if not individual:
+        final_loss = loss
+        individual_losses = []
+    else:
+        final_loss = loss[0]
+        individual_losses = loss[1]
+        
+    final_loss.backward()
+    optimizer.step()
+
+    return final_loss, individual_losses
+
+
+def train_an_epoch(
+    base_dir: str,
+    model,
+    data_loader,
+    epoch_id: int,
+    loss_func,
+    optimizer,
+    device,
+    num_warm_epochs: int = 0,
+    individual: bool = True,
+    **kwargs
+):
+    if epoch_id < num_warm_epochs:
+        model.freeze_encoder = True
+    else:
+        model.freeze_encoder = False
+
+    for batch in tqdm(data_loader):
+        loss, individual_losses = train_a_batch(
+            model=model,
+            batch_data=batch,
+            loss_func=loss_func,
+            optimizer=optimizer,
+            device=device,
+            individual=individual,
+            **kwargs
+        )
+        with open(
+            osp.join(base_dir, 'loss.log'),
+            'a'
+        ) as f:
+            f.write(str(loss.item()))
+            f.write('\t')
+            for individual_loss in individual_losses:
+                f.write(str(individual_loss))
+                f.write('\t')
+            f.write('\n')
+ 
+    ckpt_file = osp.join(
+        base_dir,
+        f'model.{epoch_id}.ckpt'
+    )
+    save_model(
+        model=model,
+        save_dir=ckpt_file,
+        epoch=epoch_id,
+        optimizer=optimizer,
+        loss=loss,
+    ) 
+
+ 
+def train_rxn(
+    base_dir,
+    model,
+    num_epochs,
+    loss_func,
+    data_loader,
+    optimizer,
+    device,
+    num_warm_epochs: int = 10,
+    ckpt_file: str = None,
+    individual: bool = True,
+    **kwargs
+):
+
+    epoch = 0
+    if ckpt_file is not None:
+
+        model, optimizer, epoch, _ = load_model(
+            ckpt_file,
+            model=model,
+            optimizer=optimizer,
+            train=True,
+            device=device,
+        )
+ 
+    for epoch_id in tnrange(num_epochs):
+
+        train_an_epoch(
+            base_dir=base_dir,
+            model=model,
+            data_loader=data_loader,
+            epoch_id=epoch + epoch_id,
+            loss_func=loss_func,
+            optimizer=optimizer,
+            num_warm_epochs=num_warm_epochs,
+            device=device,
+            individual=individual,
+            **kwargs
+        )
+
+
+def rxn_engine(
+    base_dir: str,
+    csv_file: str,
+    splitter: str,
+    smiles_col: str,
+    hard: bool = False,
+    num_epochs: int = 20,
+    target_cols: t.List = [],
+    nums_classes: t.List = [],
+    loss_func: str = 'ce',
+    num_warm_epochs: int = 10,
+    batch_size: int = 128,
+    hidden_size: int = 128,
+    lr: float = 0.01,
+    num_hidden_layers: int = 10,
+    shuffle: bool = True,
+    num_workers: int = 12,
+    dim=-1,
+    out_act='softmax',
+    device_id: int = 0,
+    individual: bool = True,
+    **kwargs
+):
+
+    base_dir = osp.abspath(base_dir)
+    makedirs(base_dir)
+    model, device = build_rxn_mu(
+        nums_classes=nums_classes,
+        hard=hard,
+        hidden_size=hidden_size,
+        nums_hidden_layers=num_hidden_layers,
+        dim=dim,
+        out_act=out_act,
+        device_id=device_id
+    )
+    if torch.cuda.device_count() > 1:
+        model = nn.DataParallel(model)
+    model.train()
+    
+    params = [{
+        'params': model.parameters(),
+        'lr': lr,
+        'weight_decay': 0
+    }]
+    optimizer = Adam(params)
+ 
+    dataset = RXNCSVDataset(
+        csv_file=csv_file,
+        splitter=splitter,
+        smiles_col=smiles_col,
+        target_cols=target_cols,
+    )
+    data_loader = RXNLoader(
+        dataset=dataset,
+        batch_size=batch_size,
+        shuffle=shuffle,
+        num_workers=num_workers
+    )
+
+    train_rxn(
+        base_dir=base_dir,
+        model=model,
+        num_epochs=num_epochs,
+        loss_func=loss_func,
+        data_loader=data_loader,
+        optimizer=optimizer,
+        device=device,
+        num_warm_epochs=num_warm_epochs,
+        ckpt_file=None,
+        individual=individual,
+        **kwargs
+    )
+
+    

hdl/controllers/train/train.py

@@ -0,0 +1,50 @@
+from torch_geometric.data import DataLoader
+import warnings
+warnings.filterwarnings("ignore")
+
+# Root mean squared error
+loss_fn = torch.nn.MSELoss()
+optimizer = torch.optim.Adam(model.parameters(), lr=0.0007)  
+
+# Use GPU for training
+device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+model = model.to(device)
+
+# Wrap data in a data loader
+data_size = len(data)
+NUM_GRAPHS_PER_BATCH = 64
+loader = DataLoader(data[:int(data_size * 0.8)], 
+                    batch_size=NUM_GRAPHS_PER_BATCH, shuffle=True)
+test_loader = DataLoader(data[int(data_size * 0.8):], 
+                         batch_size=NUM_GRAPHS_PER_BATCH, shuffle=True)
+
+def train(data):
+    # Enumerate over the data
+    for batch in loader:
+      # Use GPU
+      batch.to(device)  
+      # Reset gradients
+      optimizer.zero_grad() 
+      # Passing the node features and the connection info
+      pred, embedding = model(batch.x.float(), batch.edge_index, batch.batch) 
+      # Calculating the loss and gradients
+      loss = torch.sqrt(loss_fn(pred, batch.y))       
+      loss.backward()  
+      # Update using the gradients
+      optimizer.step()   
+    return loss, embedding
+
+print("Starting training...")
+losses = []
+for epoch in range(2000):
+    loss, h = train(data)
+    losses.append(loss)
+    if epoch % 100 == 0:
+      print(f"Epoch {epoch} | Train Loss {loss}")
+# Visualize learning (training loss)
+import seaborn as sns
+losses_float = [float(loss.cpu().detach().numpy()) for loss in losses] 
+loss_indices = [i for i,l in enumerate(losses_float)] 
+plt = sns.lineplot(loss_indices, losses_float)
+plt
+As result we get something like this:

hdl/controllers/train/train_ginet.py

@@ -0,0 +1,316 @@
+import typing as t
+from os import path as osp
+# from os import path as osp
+from itertools import cycle
+# import datetime
+
+import torch
+import numpy as np
+import pandas as pd
+
+
+# from jupyfuncs.glob import makedirs
+from jupyfuncs.pbar import tnrange, tqdm
+# from hdl.data.dataset.graph.gin import MoleculeDataset
+from hdl.data.dataset.graph.gin import MoleculeDatasetWrapper
+# from hdl.metric_loss.loss import get_lossfunc
+# from hdl.models.utils import save_model
+from .trainer_base import TorchTrainer
+
+
+class GINTrainer(TorchTrainer):
+    def __init__(
+        self,
+        base_dir,
+        data_loader,
+        test_loader,
+        metrics: t.List[str] = ['rsquared', 'rmse', 'mae'],
+        loss_func: str = 'mse',
+        model=None,
+        model_name=None,
+        model_init_args=None,
+        ckpt_file=None,
+        model_ckpt=None,
+        fix_emb=True,
+        optimizer=None,
+        optimizer_name=None,
+        optimizer_kwargs=None,
+        device=torch.device('cpu'),
+        # logger=None
+    ) -> None:
+        super().__init__(
+            base_dir=base_dir,
+            data_loader=data_loader,
+            test_loader=test_loader,
+            metrics=metrics,
+            loss_func=loss_func,
+            model=model,
+            model_name=model_name,
+            model_init_args=model_init_args,
+            ckpt_file=ckpt_file,
+            model_ckpt=model_ckpt,
+            optimizer=optimizer,
+            optimizer_name=optimizer_name,
+            optimizer_kwargs=optimizer_kwargs,
+            device=device,
+        )
+        # self.loss_func = get_lossfunc(self.loss_func)
+        # self.metrics = [get_metric(metric) for metric in metrics]
+        if fix_emb:
+            for gin in self.model.gins:
+                for param in gin.parameters():
+                    param.requires_grad = False
+    
+    def train_a_batch(self, data):
+        self.optimizer.zero_grad()
+        for i in data[: -1]:
+            for j in i:
+                j.to(self.device)
+        y = data[-1].to(self.device)
+        y = y / 100
+
+        y_pred = self.model(data).flatten()
+        
+        loss = self.loss_func(y_pred, y)
+
+        loss.backward()
+        self.optimizer.step()
+ 
+        return loss
+    
+    def load_ckpt(self):
+        self.model.load_ckpt()
+        
+    def train_an_epoch(
+        self,
+    ):
+        for i, (data, test_data) in enumerate(
+            zip(
+                self.data_loader,
+                cycle(self.test_loader)
+            )
+        ):
+            loss = self.train_a_batch(data)
+            self.losses.append(loss.item())
+            self.n_iter += 1
+            self.logger.add_scalar(
+                'train_loss',
+                loss.item(),
+                global_step=self.n_iter
+            )
+
+            if self.n_iter % 10 == 0:
+                for i in test_data[: -1]:
+                    for j in i:
+                        j.to(self.device)
+                y = test_data[-1].to(self.device)
+                y = y / 100
+
+                y_pred = self.model(test_data).flatten()
+                valid_loss = self.loss_func(y_pred, y)
+
+                y_pred = y_pred.cpu().detach().numpy()
+                y = y.cpu().detach().numpy()
+                
+                self.logger.add_scalar(
+                    'valid_loss',
+                    valid_loss.item(),
+                    global_step=self.n_iter
+                )
+
+                for metric_name, metric in zip(
+                    self.metric_names,
+                    self.metrics
+                ):
+                    self.logger.add_scalar(
+                        metric_name,
+                        metric(y_pred, y),
+                        global_step=self.n_iter
+                    )
+ 
+        self.save() 
+        self.epoch_id += 1
+    
+    def train(self, num_epochs):
+        # dir_name = datetime.now().strftime('%b%d_%H-%M-%S')
+        # makedirs(osp.join(self.base_dir, dir_name))
+
+        for _ in tnrange(num_epochs):
+            self.train_an_epoch()
+    
+    def predict(self, data_loader):
+        result_list = []
+        for data in tqdm(data_loader):
+            for i in data[: -1]:
+                for j in i:
+                    j.to(self.device)
+            # print(data[0][0].x.device)
+            # for param in self.model.parameters():
+            #     print(param.device)
+            #     break
+            y_pred = self.model(data).flatten()
+            result_list.append(y_pred.cpu().detach().numpy())
+        results = np.hstack(result_list)
+        return results
+ 
+
+def engine(
+    base_dir,
+    data_path,
+    test_data_path,
+    batch_size=128,
+    num_workers=64,
+    model_name='GINMLPR',
+    num_layers=5,
+    emb_dim=300,
+    feat_dim=512,
+    out_dim=1,
+    drop_ratio=0.0,
+    pool='mean',
+    ckpt_file=None,
+    fix_emb: bool = False,
+    device='cuda:1',
+    num_epochs=300,
+    optimizer_name='adam',
+    lr=0.001,
+    file_type: str = 'csv',
+    smiles_col_names: t.List = [],
+    y_col_name: str = None,  # "yield (%)",
+    loss_func: str = 'mse',
+    metrics: t.List[str] = ['rsquared', 'rmse', 'mae'],
+):
+    model_init_args = {
+        "num_layer": num_layers,
+        "emb_dim": emb_dim,
+        "feat_dim": feat_dim,
+        "out_dim": out_dim,
+        "drop_ratio": drop_ratio,
+        "pool": pool,
+        "ckpt_file": ckpt_file,
+        "num_smiles": len(smiles_col_names),
+    }
+    wrapper = MoleculeDatasetWrapper(
+        batch_size=batch_size,
+        num_workers=num_workers,
+        valid_size=0,
+        data_path=data_path,
+        file_type=file_type,
+        smi_col_names=smiles_col_names,
+        y_col_name=y_col_name
+    )
+    test_wrapper = MoleculeDatasetWrapper(
+        batch_size=batch_size,
+        num_workers=num_workers,
+        valid_size=0,
+        data_path=test_data_path,
+        file_type=file_type,
+        smi_col_names=smiles_col_names,
+        y_col_name=y_col_name
+    )
+
+    data_loader = wrapper.get_test_loader(
+        shuffle=True
+    )
+    test_loader = test_wrapper.get_test_loader(
+        shuffle=False
+    ) 
+
+    trainer = GINTrainer(
+        base_dir=base_dir,
+        model_name=model_name,
+        model_init_args=model_init_args,
+        optimizer_name=optimizer_name,
+        ckpt_file=ckpt_file,
+        fix_emb=fix_emb,
+        optimizer_kwargs={"lr": lr},
+        data_loader=data_loader,
+        test_loader=test_loader,
+        metrics=metrics,
+        loss_func=loss_func,
+        device=device
+    )
+    
+    trainer.train(num_epochs=num_epochs)
+
+
+def predict(
+    base_dir,
+    data_path,
+    batch_size=128,
+    num_workers=64,
+    model_name='GINMLPR',
+    num_layers=5,
+    emb_dim=300,
+    feat_dim=512,
+    out_dim=1,
+    drop_ratio=0.0,
+    pool='mean',
+    ckpt_file=None,
+    model_ckpt=None,
+    device='cuda:1',
+    file_type: str = 'csv',
+    smiles_col_names: t.List = [],
+    y_col_name: str = None,  # "yield (%)",
+    metrics: t.List[str] = ['rsquared', 'rmse', 'mae'],
+):
+    model_init_args = {
+        "num_layer": num_layers,
+        "emb_dim": emb_dim,
+        "feat_dim": feat_dim,
+        "out_dim": out_dim,
+        "drop_ratio": drop_ratio,
+        "pool": pool,
+        "ckpt_file": ckpt_file,
+        "num_smiles": len(smiles_col_names),
+    }
+    wrapper = MoleculeDatasetWrapper(
+        batch_size=batch_size,
+        num_workers=num_workers,
+        valid_size=0,
+        data_path=data_path,
+        file_type=file_type,
+        smi_col_names=smiles_col_names,
+        y_col_name=y_col_name
+    )
+    data_loader = wrapper.get_test_loader(
+        shuffle=False
+    )
+    trainer = GINTrainer(
+        base_dir=base_dir,
+        model_name=model_name,
+        model_init_args=model_init_args,
+        model_ckpt=model_ckpt,
+        data_loader=data_loader,
+        test_loader=None,
+        metrics=metrics,
+        device=device
+    )
+    metric_list = trainer.metrics
+    trainer.load(ckpt_file=model_ckpt)
+    trainer.model.eval()
+    results = trainer.predict(data_loader)
+
+    df = pd.read_csv(data_path)
+    df['pred'] = results
+    df.to_csv(
+        osp.join(base_dir, 'pred.csv'),
+        index=False
+    )
+    
+    if y_col_name is not None:
+        metrics_df = pd.DataFrame()
+        y = df[y_col_name].array / 100
+        for metric_name, metric in zip(
+            metrics,
+            metric_list
+        ):
+            metrics_df[metric_name] = np.array([metric(
+                y, results
+            )])
+        metrics_df.to_csv(
+            osp.join(
+                base_dir, 'metrics.csv'
+            ),
+            index=False
+        )
+