update

Author: AidenSb

apamodel/blocks.py

@@ -64,3 +64,36 @@ def __init__(self, layer_dims, dropouts, dropout=False):
 
     def forward(self, x):
         return self.op(x)
+
+
+class ProcessSelfAttn(nn.Module):
+    """
+    Implements the self-attention mechanism.
+    Attributes
+    ----------
+    nhead : int
+        The number of attention heads.
+        Each head computes a separate attention score for each token.
+    """
+
+    def __init__(
+        self,
+        embed_dim: int,
+        num_layers: int,
+        nhead: int,
+        dim_feedforward: int = 2048,
+        dropout: float = 0.2,
+    ):
+        super().__init__()
+        self.encoder_layer = nn.TransformerEncoderLayer(
+            embed_dim,
+            nhead,
+            dim_feedforward,
+            dropout,
+            activation="gelu",
+            batch_first=True,
+        )
+        self.transformer = nn.TransformerEncoder(self.encoder_layer, num_layers)
+
+    def forward(self, latent):
+        return self.transformer(latent)

apamodel/model.py

@@ -2,13 +2,14 @@
 import torch.nn as nn
 import torch.optim as optim
 from torch.utils.data import Dataset, DataLoader
-from blocks import ConvBlock, FCBlock
+from blocks import ConvBlock, FCBlock, ProcessSelfAttn
 import numpy as np
 
-RBP_COUNT = 279
+RBP_COUNT = 327
 FIX_SEQ_LEN = 4000
 
 
+
 class APAData(Dataset):
     """
     APAData is a dataset class for APA-Net model.
@@ -19,40 +20,30 @@ class APAData(Dataset):
         device (str): Device to use (e.g., 'cuda' or 'cpu').
     """
 
-    def __init__(self, seqs, df, ct, device):
+    def __init__(self, data, device):
         self.device = device
         self.reg_label = torch.from_numpy(
-            np.array(df[:, 3].tolist(), dtype=np.float32)
+            np.array(data[:, 7].tolist(), dtype=np.float32)
         ).to(device)
-        self.seq_idx = torch.from_numpy(np.array(df[:, 1].tolist(), dtype=np.int32)).to(
+        self.oneH_seqs = torch.from_numpy(np.array(data[:, 6].tolist())).to(
             device
         )
-        self.oneH_seqs = torch.from_numpy(np.array(list(seqs[:, 3]), dtype=np.int8)).to(
+        self.ct_profiles = torch.from_numpy(np.array(data[:, 8].tolist())).to(
             device
         )
-        self.oneH_seq_indexes = torch.from_numpy(
-            np.array(seqs[:, 0], dtype=np.int32)
-        ).to(device)
-        self.celltypes = df[:, 2]
-        self.ct_profiles = ct
+        self.celltype_name = data[:, 1].tolist()
+        self.switch_name = data[:, 5].tolist()
 
     def __len__(self):
         return self.reg_label.shape[0]
 
     def __getitem__(self, idx):
-        seq_idx = self.seq_idx[idx]
-        seq = (
-            self.oneH_seqs[torch.where(self.oneH_seq_indexes == seq_idx)]
-            .squeeze()
-            .type(torch.cuda.FloatTensor)
-        )
+        seq = self.oneH_seqs[idx].type(torch.cuda.FloatTensor)
         reg_label = self.reg_label[idx]
-        celltype_name = self.celltypes[idx]
-        celltype = torch.from_numpy(
-            self.ct_profiles[celltype_name].values.astype(np.float32)
-        ).to(self.device)
-        return (seq, celltype, celltype_name, reg_label)
-
+        celltype_profile = self.ct_profiles[idx].type(torch.cuda.FloatTensor)
+        celltype_name = self.celltype_name[idx]
+        switch_name = self.switch_name[idx]
+        return (seq, reg_label, celltype_profile, celltype_name, switch_name)
 
 class APANET(nn.Module):
     """
@@ -108,6 +99,13 @@ def _build_model(self):
             dropouts=self.config["fc2_dropouts"],
             dropout=True,
         )
+        self.process_self_attn = ProcessSelfAttn(
+            self.config["psa_query_dim"],
+            self.config["psa_num_layers"],
+            self.config["psa_nhead"],
+            self.config["psa_dim_feedforward"],
+            self.config["psa_dropout"]
+        )
 
     def _get_conv1d_out_length(self, l_in, kernel, stride, pool_kernel, pool_stride):
         """Utility method to calculate output length of Conv1D layer."""
@@ -118,10 +116,11 @@ def _get_conv1d_out_length(self, l_in, kernel, stride, pool_kernel, pool_stride)
 
     def forward(self, seq, celltype):
         # Convolutional forward
-        x_conv = self.conv_block_1(seq)
-        x = x_conv.permute(2, 0, 1)  # reshape for attention block
-        x, _ = self.attention(x, x, x)
-        x = x.permute(1, 2, 0)  # reshape back
+        x_conv = self.conv_block_1(seq) # batch, 64/128(dim), 80(len)
+        x = x_conv.permute(0, 2, 1)  # reshape for attention block so dim is first
+        # x, _ = self.attention(x, x, x)
+        x = self.process_self_attn(x)
+        x = x.permute(0, 2, 1)  # reshape back
         x = x + x_conv  # add residual connection
         x = torch.flatten(x, 1)  # flatten for FC layers
         x = self.fc1(x)  # FC block 1

apamodel/train_script.py

@@ -11,7 +11,7 @@
 
 
 def build_dataloaders(
-    device, train_seq, valid_seq, train_data, val_data, batch_size, ct_profiles
+    device, train_data, valid_data, batch_size,
 ):
     """
     Create training and validation data loaders.
@@ -24,21 +24,27 @@ def build_dataloaders(
         Tuple of DataLoader for training and validation datasets.
     """
     train_loader = DataLoader(
-        APAData(train_seq, train_data, ct_profiles, device),
+        APAData(train_data, device),
         batch_size=batch_size,
         shuffle=True,
         drop_last=True,
     )
     valid_loader = DataLoader(
-        APAData(valid_seq, val_data, ct_profiles, device),
+        APAData(valid_data, device),
         batch_size=batch_size,
         shuffle=False,
         drop_last=False,
     )
     return train_loader, valid_loader
 
+def l1_penalty(model, l1_factor):
+    l1_reg = torch.tensor(0.).to(model.device)
+    for param in model.parameters():
+        l1_reg += torch.norm(param, 1)
+    return l1_factor * l1_reg
 
-def train_one_epoch(model, train_loader):
+
+def train_one_epoch(model, train_loader, l1_factor=0.00005):
     """
     Train the model for one epoch.
     Args:
@@ -49,10 +55,13 @@ def train_one_epoch(model, train_loader):
     """
     model.train()
     total_loss, predictions, targets = 0.0, [], []
-    for seq_X, celltype, _, Y in train_loader:
+    for seq_X, Y, celltype, _, _ in train_loader:
         model.optimizer.zero_grad()
         outputs = torch.squeeze(model(seq_X, celltype))
-        loss = torch.sqrt(model.loss_fn(outputs, Y))
+        mse_loss = torch.sqrt(model.loss_fn(outputs, Y))
+        # l1_loss = l1_penalty(model, l1_factor)
+        # loss = mse_loss + l1_loss 
+        loss = mse_loss
         loss.backward()
         model.optimizer.step()
         total_loss += loss.item() * seq_X.size(0)
@@ -77,7 +86,7 @@ def validate_one_epoch(model, valid_loader):
     model.eval()
     total_loss, predictions, targets = 0.0, [], []
     with torch.no_grad():
-        for seq_X, celltype, _, Y in valid_loader:
+        for seq_X, Y, celltype, _, _ in valid_loader:
             outputs = torch.squeeze(model(seq_X, celltype))
             loss = torch.sqrt(model.loss_fn(outputs, Y))
             total_loss += loss.item() * seq_X.size(0)
@@ -92,13 +101,11 @@ def validate_one_epoch(model, valid_loader):
 
 
 def main_train(
-    train_seq,
-    valid_seq,
     train_data,
     val_data,
-    profiles,
     modelfile,
     device,
+    project_name,
     config,
     use_wandb,
 ):
@@ -115,18 +122,15 @@ def main_train(
     use_wandb = args.use_wandb.lower() == "true"
     train_loader, valid_loader = build_dataloaders(
         device,
-        train_seq,
-        valid_seq,
         train_data,
         val_data,
         config["batch_size"],
-        profiles,
     )
     with tqdm(range(config["epochs"]), unit="epoch") as tepochs:
         if use_wandb:
             wandb.login()
             with wandb.init(
-                project=config["project_name"],
+                project= project_name,
                 settings=wandb.Settings(start_method="thread"),
             ):
                 model = APANET(config)
@@ -167,18 +171,9 @@ def main_train(
     parser.add_argument(
         "--train_data", type=str, required=True, help="Path to training data file"
     )
-    parser.add_argument(
-        "--train_seq", type=str, required=True, help="Path to training sequences file"
-    )
     parser.add_argument(
         "--valid_data", type=str, required=True, help="Path to validation data file"
     )
-    parser.add_argument(
-        "--valid_seq", type=str, required=True, help="Path to validation sequences file"
-    )
-    parser.add_argument(
-        "--profiles", type=str, required=True, help="Path to cell type profiles file"
-    )
     parser.add_argument(
         "--modelfile", type=str, required=True, help="Path to save the trained model"
     )
@@ -214,10 +209,7 @@ def main_train(
     np.random.seed(7)
 
     train_data = np.load(args.train_data, allow_pickle=True)
-    train_seq = np.load(args.train_seq, allow_pickle=True)
     valid_data = np.load(args.valid_data, allow_pickle=True)
-    valid_seq = np.load(args.valid_seq, allow_pickle=True)
-    profiles = pd.read_csv(args.profiles, index_col=0, sep="\t")
 
     config = {
         "batch_size": args.batch_size,
@@ -227,35 +219,38 @@ def main_train(
         "opt": "Adam",
         "loss": "mse",
         "lr": 2.5e-05,
-        "adam_weight_decay": 0.06,
-        "conv1kc": 128,
+        "adam_weight_decay": 0.09, # 0.06 before
+        "conv1kc": 128, #128, 64
         "conv1ks": 12,
         "conv1st": 1,
-        "pool1ks": 25,
-        "pool1st": 25,
-        "cnvpdrop1": 0.2,
+        "pool1ks": 16,
+        "pool1st": 16,
+        "cnvpdrop1": 0,
         "Matt_heads": 8,
         "Matt_drop": 0.2,
         "fc1_dims": [
-            8192,
+            8192, # 8192, 5120
             4048,
             1024,
             512,
             256,
         ],  # first dimension will be calculated dynamically
-        "fc1_dropouts": [0.3, 0.25, 0.25, 0.2, 0.1],
+        "fc1_dropouts": [0.25, 0.25, 0.25, 0, 0],
         "fc2_dims": [128, 32, 16, 1],  # first dimension will be calculated dynamically
         "fc2_dropouts": [0.2, 0.2, 0, 0],
+        'psa_query_dim': 128, # make sure this is correct
+        'psa_num_layers': 1,
+        'psa_nhead': 1,
+        'psa_dim_feedforward':1024,
+        'psa_dropout': 0 
     }
 
     main_train(
-        train_seq,
-        valid_seq,
         train_data,
         valid_data,
-        profiles,
         args.modelfile,
         args.device,
+        args.project_name,
         config,
         args.use_wandb,
     )