sync android mnist binary loading;fix ios data type

FedCampus · Sep 25, 2023 · afd7f49 · afd7f49
1 parent cf28ea2
commit afd7f49
Show file tree

Hide file tree

Showing 5 changed files with 12 additions and 11 deletions.
diff --git a/android/fed_kit_examples/src/main/java/org/eu/fedcampus/fed_kit_examples/mnist/mnist.kt b/android/fed_kit_examples/src/main/java/org/eu/fedcampus/fed_kit_examples/mnist/mnist.kt
@@ -14,7 +14,7 @@ import java.io.File
 fun sampleSpec() = SampleSpec<Float3DArray, FloatArray>(
     { it.toTypedArray() },
     { it.toTypedArray() },
-    { Array(it) { FloatArray(N_CLASSES) } },
+    { Array(it) { FloatArray(1) } },
     ::maxSquaredErrorLoss,
     { samples, logits ->
         samples.zip(logits).forEach { (sample, logit) ->
@@ -60,16 +60,17 @@ private fun addSample(
 ) {
     val splits = line.split(",")
     val feature = Array(IMAGE_SIZE) { Array(IMAGE_SIZE) { FloatArray(1) } }
-    val label = FloatArray(N_CLASSES)
+    val label = FloatArray(1)
     for (i in 0 until LENGTH_ENTRY) {
         feature[i / IMAGE_SIZE][i % IMAGE_SIZE][0] = splits[i + 1].toFloat() / NORMALIZATION
     }
-    label[splits.first().toInt()] = 1f
+    if (splits.first().toInt() == 1) {
+        label[0] = 1f
+    }
     flowerClient.addSample(feature, label, isTraining)
 }
 
 private const val TAG = "MNIST Data Loader"
 private const val IMAGE_SIZE = 28
-private const val N_CLASSES = 10
 private const val LENGTH_ENTRY = IMAGE_SIZE * IMAGE_SIZE
 private const val NORMALIZATION = 255f
diff --git a/backend/train/run.py b/backend/train/run.py
@@ -69,7 +69,7 @@ def fit_config(_: int):
     """
     config: dict[str, Scalar] = {
         "batch_size": 32,
-        "local_epochs": 2,
+        "local_epochs": 10,
     }
     return config
 

diff --git a/fed_kit_client/ios/Runner/DataLoader.swift b/fed_kit_client/ios/Runner/DataLoader.swift
@@ -122,14 +122,14 @@ private func prepareMLBatchProvider(
             group.addTask {
                 let splits = line.split(separator: ",")
                 let imageMultiArr = try! MLMultiArray(shape: shapeData, dataType: .float32)
-                let outputMultiArr = try! MLMultiArray(shape: shapeTarget, dataType: .int32)
+                let outputMultiArr = try! MLMultiArray(shape: shapeTarget, dataType: .double)
                 for i in 0 ..< lengthEntry {
                     imageMultiArr[i] = (Float(String(splits[i + 1]))! / normalization) as NSNumber
                 }
                 if Int(String(splits[0]))! == 1 {
-                    outputMultiArr[0] = 1
+                    outputMultiArr[0] = 1.0
                 } else {
-                    outputMultiArr[0] = 0
+                    outputMultiArr[0] = 0.0
                 }
                 let imageValue = MLFeatureValue(multiArray: imageMultiArr)
                 let outputValue = MLFeatureValue(multiArray: outputMultiArr)

diff --git a/fed_kit_client/ios/Runner/Utils/Flwr/MLClient.swift b/fed_kit_client/ios/Runner/Utils/Flwr/MLClient.swift
@@ -103,9 +103,9 @@ public class MLClient {
             }
             let actl = batch.features(at: index).featureValue(for: modelProto.target)!
             let prediction = try pred.multiArrayValue!.toArray(type: Float.self)
-            let actual = try actl.multiArrayValue!.toArray(type: Int32.self)
+            let actual = try actl.multiArrayValue!.toArray(type: Double.self)
             totalLoss += meanSquareErrors(prediction, actual)
-            if actual.argmax() == prediction.argmax() {
+            if (actual[0] == 0 && prediction[0] < 0.5) || (actual[0] != 0 && prediction[0] >= 0.5) {
                 nCorrect += 1
             }
         }

diff --git a/fed_kit_client/ios/Runner/Utils/MLHelper.swift b/fed_kit_client/ios/Runner/Utils/MLHelper.swift
@@ -70,7 +70,7 @@ struct Layer {
 }
 
 /// Assuming the two have the same length.
-func meanSquareErrors(_ a: [Float], _ b: [Int32]) -> Float {
+func meanSquareErrors(_ a: [Float], _ b: [Double]) -> Float {
     var sum: Float = 0.0
     for (index, value) in a.enumerated() {
         let diff = value - Float(b[index])