update for the model and loss function

ddi/dataset.py

@@ -58,7 +58,6 @@ def construct_load_dataloaders(dataset_fold, dsettypes, config, wrk_dir):
     # setup data loaders
     data_loaders = {}
     epoch_loss_avgbatch = {}
-    epoch_loss_avgsamples = {}
     flog_out = {}
     score_dict = {}
     class_weights = {}
@@ -75,19 +74,19 @@ def construct_load_dataloaders(dataset_fold, dsettypes, config, wrk_dir):
                                             num_workers=config['num_workers'])
 
         epoch_loss_avgbatch[dsettype] = []
-        epoch_loss_avgsamples[dsettype] = []
         score_dict[dsettype] = ModelScore(0, 0.0, 0.0, 0.0, 0.0, 0.0)  # (best_epoch, auc, aupr, f1, precision, recall)
         if(wrk_dir):
             flog_out[dsettype] = os.path.join(wrk_dir, dsettype + ".log")
         else:
             flog_out[dsettype] = None
 
-    return (data_loaders, epoch_loss_avgbatch, epoch_loss_avgsamples, score_dict, class_weights, flog_out)
+    return (data_loaders, epoch_loss_avgbatch, score_dict, class_weights, flog_out)
 
 def preprocess_features(feat_fpath):
     X_fea = np.loadtxt(feat_fpath,dtype=float,delimiter=",")
     r, c = np.triu_indices(len(X_fea),1) # take indices off the diagnoal by 1
     return np.concatenate((X_fea[r], X_fea[c]), axis=1)
+
 def preprocess_labels(interaction_fpath):
     interaction_matrix = np.loadtxt(interaction_fpath,dtype=float,delimiter=",")
     r, c = np.triu_indices(len(interaction_matrix),1) # take indices off the diagnoal by 1
@@ -102,7 +101,7 @@ def create_setvector_features(X, num_sim_types):
     h = np.transpose(g, axes=(2,0, 1))
     return h
 
-def get_stratified_partitions(ddi_datatensor, num_folds=5, valid_set_portion=0.1, random_state=42):
+def get_stratified_partitions(ddi_datatensor, num_folds=5, random_state=42):
     """Generate 5-fold stratified sample of drug-pair ids based on the interaction label
 
     Args:
@@ -127,6 +126,32 @@ def get_stratified_partitions(ddi_datatensor, num_folds=5, valid_set_portion=0.1
         print("-"*25)
     return(data_partitions)
 
+def get_validation_partitions(ddi_datatensor, num_folds=2, valid_set_portion=0.1, random_state=42):
+    """Generate stratified train/validation split of drug-pair ids based on the interaction label
+
+    Args:
+        ddi_datatensor: instance of :class:`DDIDataTensor`
+    """
+    skf_trte = StratifiedShuffleSplit(n_splits=num_folds, 
+                                      test_size=valid_set_portion, 
+                                      random_state=random_state)  # split train and test
+    data_partitions = {}
+    X = ddi_datatensor.X_feat
+    y = ddi_datatensor.y
+    fold_num = 0
+    for train_index, test_index in skf_trte.split(X,y):
+
+        data_partitions[fold_num] = {'train': train_index,
+                                     'validation': test_index}
+        print("fold_num:", fold_num)
+        print('train data')
+        report_label_distrib(y[train_index])
+        print('validation data')
+        report_label_distrib(y[test_index])
+        print()
+        fold_num += 1
+        print("-"*25)
+    return(data_partitions)
 
 def report_label_distrib(labels):
     classes, counts = np.unique(labels, return_counts=True)
@@ -168,6 +193,8 @@ def generate_partition_datatensor(ddi_datatensor, data_partitions):
             target_ids = data_partitions[fold_num][dsettype]
             datatensor_partition = PartitionDataTensor(ddi_datatensor, target_ids, dsettype, fold_num)
             datatensor_partitions[fold_num][dsettype] = datatensor_partition
+    compute_class_weights_per_fold_(datatensor_partitions)
+
     return(datatensor_partitions)
 
 def build_datatensor_partitions(data_partitions, ddi_datatensor):

ddi/model_attn.py

@@ -108,10 +108,51 @@ def forward(self, X):
 
         return z
 
+class FeatureEmbAttention(nn.Module):
+    def __init__(self, input_dim):
+        '''
+        Args:
+            input_dim: int, size of the input vector (i.e. feature vector)
+        '''
+
+        super().__init__()
+        self.input_dim = input_dim
+        # use this as query vector against the transformer outputs
+        self.queryv = nn.Parameter(torch.randn(input_dim, dtype=torch.float32), requires_grad=True)
+        self.softmax = nn.Softmax(dim=1) # normalized across seqlen
+
+    def forward(self, X):
+        '''Performs forward computation
+        Args:
+            X: torch.Tensor, (batch, ddi similarity type vector, feat_dim), dtype=torch.float32
+        '''
+
+        X_scaled = X / (self.input_dim ** (1/4))
+        queryv_scaled = self.queryv / (self.input_dim ** (1/4))
+        # using  matmul to compute tensor vector multiplication
+
+        # (bsize, seqlen)
+        attn_weights = X_scaled.matmul(queryv_scaled)
+
+        # softmax
+        attn_weights_norm = self.softmax(attn_weights)
+
+        # reweighted value vectors (in this case reweighting the original input X)
+        # unsqueeze attn_weights_norm to get (bsize, 1, num similarity type vectors)
+        # perform batch multiplication with X that has shape (bsize, num similarity type vectors, feat_dim)
+        # result will be (bsize, 1, feat_dim)
+        # squeeze the result to obtain (bsize, feat_dim)
+        z = attn_weights_norm.unsqueeze(1).bmm(X).squeeze(1)
+
+        # returns (bsize, feat_dim), (bsize, num similarity type vectors)
+        return z, attn_weights_norm
+
+
 class DDI_Transformer(nn.Module):
 
     def __init__(self, input_size=586, num_attn_heads=8, mlp_embed_factor=2, 
-                nonlin_func=nn.ReLU(), pdropout=0.3, num_transformer_units=12, num_classes=1):
+                nonlin_func=nn.ReLU(), pdropout=0.3, num_transformer_units=12,
+                pooling_mode = 'attn', num_classes=2):
 
         super(DDI_Transformer, self).__init__()
 
@@ -122,9 +163,25 @@ def __init__(self, input_size=586, num_attn_heads=8, mlp_embed_factor=2,
         self.trfunit_pipeline = nn.Sequential(*trfunit_layers)
         embed_size = input_size
         self.Wy = nn.Linear(embed_size, num_classes)
+        self.pooling_mode = pooling_mode
+        if pooling_mode == 'attn':
+            self.pooling = FeatureEmbAttention(input_size)
+        elif pooling_mode == 'mean':
+            self.pooling = torch.mean
+
         # perform log softmax on the feature dimension
-        # self.log_softmax = nn.LogSoftmax(dim=-1)
-
+        self.log_softmax = nn.LogSoftmax(dim=-1)
+        self._init_params_()
+
+
+    def _init_params_(self):
+        for p_name, p in self.named_parameters():
+            param_dim = p.dim()
+            if param_dim > 1: # weight matrices
+                nn.init.xavier_uniform_(p)
+            elif param_dim == 1: # bias parameters
+                if p_name.endswith('bias'):
+                    nn.init.uniform_(p, a=-1.0, b=1.0)
 
     def forward(self, X):
         """
@@ -137,7 +194,16 @@ def forward(self, X):
         # pool across similarity type vectors
         # Note: z.mean(dim=1) will change shape of z to become (batch, input_size)
         # we can keep dimension by running z.mean(dim=1, keepdim=True) to have (batch, 1, input_size)
+
+        # pool across similarity type vectors
+        if self.pooling_mode == 'attn':
+            z, fattn_w_norm = self.pooling(z)
+        # Note: z.mean(dim=1) or self.pooling(z, dim=1) will change shape of z to become (batch, embedding dim)
+        # we can keep dimension by running z.mean(dim=1, keepdim=True) to have (batch, 1, embedding dim)
+        elif self.pooling_mode == 'mean':
+            z = self.pooling(z, dim=1)
+            fattn_w_norm = None
+
+        y = self.Wy(z) 
 
-        y = self.Wy(z.mean(dim=1)) 
-
-        return y
+        return self.log_softmax(y), fattn_w_norm