Federated Learning System

federated learning Training a model across several devices that each keep their own data, sending only the learned weights to a central server. is a method where the training data stays on the local machines and only the learned weights are sent to a central server. This project shows the benefits and the performance of that approach.

It is good for performance because you do not move all the data to one machine to train it. The same reason makes it good for privacy and security, since the raw data never leaves each device.

Each Raspberry Pi, using Camera Module 3, takes photos of metal parts to grow its local dataset. It trains or predicts locally, then its weights are aggregated on the central server for a stronger model. The server only ever holds weights, never the data, so the data stays private on each device.

A goal was a secure and stable link between the clients and the server. I used SSH for that and transferred the weights in under 6 seconds per client.

The dataset had 6666 metal images, split for 2-client and 3-client runs, each with training, validation, and testing, across 4 defect classes.

The first run over-fit on class 3 (scratches). That pushed me to clean the data properly and use balanced sampling. F1 score went up by 8.51 percent and accuracy by 10.10 percent.

What do these curves mean?

A PR curve shows how well the model finds positives without making many mistakes. The first curve sits near the bottom left, so it makes mistakes as it predicts more positives. The second moves toward the top right, which is good: more positives without many more mistakes.

To measure the gain from more devices, I ran 2 clients and 3 clients and compared them.

Going from 2 to 3 clients added 5.39 percent F1 score and 6.56 percent accuracy, which shows how the approach scales with more devices.

I used PyTorch with YOLOv8 on the desktops and TensorFlow Lite on the Pis. Only 5 epochs were possible due to hardware, and up to 3 clients. Even so the model reached close to 60 percent, where federated models usually land at 70 to 75 percent, so there is clear room to grow with more epochs and devices.

The detection here was metal defects, but after talking with a doctor friend I want to take it to radiograph analysis. Federated training would let hospitals train locally and keep patient data private.