add metrics for multilabel problems

Orion-AI-Lab · Jun 17, 2023 · 294323f · 294323f
1 parent fece70f
commit 294323f
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 8 deletions.
diff --git a/configs/supervised_configs.json b/configs/supervised_configs.json
@@ -1,10 +1,10 @@
 {
-    "wandb_project":"YOUR_PROJECT",
+    "wandb_project":"YOUR_WANDB_PROJECT",
     "wandb_entity":"YOUR_ENTITY",
     "task":"classification",
     "num_classes":11,
     "device":"cuda:1",
-    "wandb":false,
+    "wandb":true,
     "mixed_precision":true,
     "ssl_encoder":null,
     "ssl_run_id_path":null,
@@ -14,9 +14,10 @@
     "num_workers":4,
     "lr":0.0001,
     "weight_decay":1e-4,
-    "epochs":10,
-    "architecture":"ResNet50",
+    "epochs":2,
+    "architecture":"ResNet18",
     "oversampling":true,
+    "multilabel":true,
     "augment":false,
     "num_channel":2,
     "class_weights":[1.0,1.0],

diff --git a/utilities/utils.py b/utilities/utils.py
@@ -118,10 +118,16 @@ def prepare_supervised_learning_loaders(configs):
     return train_loader, val_loader, test_loader
 
 def initialize_metrics(configs):
-    accuracy = Accuracy(task='multiclass', average='micro',multidim_average='global',num_classes=configs['num_classes']).to(configs['device'])
-    fscore = F1Score(task='multiclass', num_classes=configs['num_classes'],average='micro',multidim_average='global').to(configs['device'])
-    precision = Precision(task='multiclass', average='micro', num_classes=configs['num_classes'],multidim_average='global').to(configs['device'])
-    recall = Recall(task='multiclass', average='micro', num_classes=configs['num_classes'],multidim_average='global').to(configs['device'])
+    if configs['multilabel']:
+        accuracy = Accuracy(task='multilabel', average='micro',multidim_average='global',num_labels=configs['num_classes']).to(configs['device'])
+        fscore = F1Score(task='multilabel', num_labels=configs['num_classes'],average='micro',multidim_average='global').to(configs['device'])
+        precision = Precision(task='multilabel', average='micro', num_labels=configs['num_classes'],multidim_average='global').to(configs['device'])
+        recall = Recall(task='multilabel', average='micro', num_labels=configs['num_classes'],multidim_average='global').to(configs['device'])
+    else:
+        accuracy = Accuracy(task='multiclass', average='micro',multidim_average='global',num_classes=configs['num_classes']).to(configs['device'])
+        fscore = F1Score(task='multiclass', num_classes=configs['num_classes'],average='micro',multidim_average='global').to(configs['device'])
+        precision = Precision(task='multiclass', average='micro', num_classes=configs['num_classes'],multidim_average='global').to(configs['device'])
+        recall = Recall(task='multiclass', average='micro', num_classes=configs['num_classes'],multidim_average='global').to(configs['device'])
     return [accuracy, fscore, precision, recall]