move weights to gpu/cpu; add stats for num of pred per sample

Author: ankitdhall

network/evaluation.py

@@ -200,6 +200,12 @@ def evaluate(self, predicted_scores, correct_labels, epoch, phase, save_to_tenso
 
         self.predicted_labels = predicted_scores >= np.tile(self.optimal_thresholds, (correct_labels.shape[0], 1))
 
+        classes_predicted_per_sample = np.sum(self.predicted_labels, axis=1)
+        print("Max: {}".format(np.max(classes_predicted_per_sample)))
+        print("Min: {}".format(np.min(classes_predicted_per_sample)))
+        print("Mean: {}".format(np.mean(classes_predicted_per_sample)))
+        print("std: {}".format(np.std(classes_predicted_per_sample)))
+
         level_stop, level_start = [], []
         for level_id, level_len in enumerate(self.labelmap.levels):
             if level_id == 0:

network/loss.py

@@ -8,7 +8,7 @@ def __init__(self, labelmap, level_weights=None, weight=None):
         torch.nn.Module.__init__(self)
         self.labelmap = labelmap
         self.level_weights = [1.0] * len(self.labelmap.levels) if level_weights is None else level_weights
-
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
         self.criterion = []
         if weight is None:
             for level_len in self.labelmap.levels:
@@ -22,7 +22,7 @@ def __init__(self, labelmap, level_weights=None, weight=None):
                 else:
                     level_start.append(level_stop[level_id - 1])
                     level_stop.append(level_stop[level_id - 1] + level_len)
-                self.criterion.append(nn.CrossEntropyLoss(weight=weight[level_start[level_id]:level_stop[level_id]],
+                self.criterion.append(nn.CrossEntropyLoss(weight=weight[level_start[level_id]:level_stop[level_id]].to(self.device),
                                                           reduction='none'))
 
         print('==Using the following weights config for multi level cross entropy loss: {}'.format(self.level_weights))
@@ -52,6 +52,9 @@ def forward(self, outputs, labels, level_labels):
 class MultiLabelSMLoss(torch.nn.MultiLabelSoftMarginLoss):
     def __init__(self, weight=None, size_average=None, reduce=None, reduction='mean'):
         print(weight)
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        if weight is not None:
+            weight = weight.to(self.device)
         torch.nn.MultiLabelSoftMarginLoss.__init__(self, weight, size_average, reduce, reduction)
 
     def forward(self, outputs, labels, level_labels):