Update loss.py

flamby/datasets/fed_synthetic/dataset_creation_scripts/download.py

@@ -134,7 +134,7 @@ def main(output_folder, debug=False, **kwargs):
         action="store_false",
         help="Generate a regression dataset. (Default)",
     )
-    parser.set_defaults(classification=False)
+    parser.set_defaults(classification=True)
 
     parser.add_argument(
         "--output-folder",

flamby/datasets/fed_synthetic/loss.py

@@ -5,7 +5,7 @@
 class BaselineLoss(_Loss):
     def __init__(self):
         super(BaselineLoss, self).__init__()
-        self.bce = torch.nn.BCELoss()
+        self.ce = torch.nn.CrossEntropyLoss()
 
     def forward(self, input: torch.Tensor, target: torch.Tensor):
-        return self.bce(input, target)
+        return self.ce(input, target.squeeze().long())

flamby/datasets/fed_synthetic/metric.py

@@ -1,14 +1,9 @@
 import numpy as np
+from sklearn.metrics import accuracy_score
 
 
 def metric(y_true, y_pred):
     y_true = y_true.astype("uint8")
-    # The try except is needed because when the metric is batched some batches
-    # have one class only
-    try:
-        # return roc_auc_score(y_true, y_pred)
-        # proposed modification in order to get a metric that calcs on center 2
-        # (y=1 only on that center)
-        return ((y_pred > 0.5) == y_true).mean()
-    except ValueError:
-        return np.nan
+    y_pred = np.argmax(y_pred, axis=1)
+    return accuracy_score(y_true, y_pred)
+

flamby/datasets/fed_synthetic/model.py

@@ -3,7 +3,7 @@
 
 
 class Baseline(nn.Module):
-    def __init__(self, input_dim=10, output_dim=1):
+    def __init__(self, input_dim=10, output_dim=3):
         super(Baseline, self).__init__()
         self.linear = torch.nn.Linear(input_dim, output_dim)