优化文档字符串，明确参数说明

fed_algo_cs/server_base.py

@@ -11,13 +11,13 @@ class FedYoloServer(object):
     def __init__(self, client_list, model_name, params):
         """
         Federated YOLO Server
-        Args:
+        Attributes:
             client_list: list of connected clients
             model_name: YOLO model architecture name
             params: dict of hyperparameters (must include 'names')
         """
         # Track client updates
-        self.client_state = {}
+        self.client_state: dict[str, dict[str, torch.Tensor]] = {}
         self.client_loss = {}
         self.client_n_data = {}
         self.selected_clients = []
@@ -64,14 +64,19 @@ class FedYoloServer(object):
                 self.selected_clients.append(client_id)
                 self.n_data += self.client_n_data[client_id]
 
+    # TODO: skip the layer which can not be learnted locally
     @torch.no_grad()
-    def agg(self):
+    def agg(self, skip_bn_layer: bool = False):
         """
         Server aggregates the local updates from selected clients using FedAvg.
 
-        :return: model_state: aggregated model weights
-        :return: avg_loss: weighted average training loss across selected clients
-        :return: n_data: total number of data points across selected clients
+        Args:
+            skip_bn_layer: whether to skip batch normalization layers during aggregation
+
+        Returns:
+            :model_state: aggregated model weights
+            :avg_loss: weighted average training loss across selected clients
+            :n_data: total number of data points across selected clients
         """
         if len(self.selected_clients) == 0 or self.n_data == 0:
             import warnings
@@ -144,11 +149,13 @@ class FedYoloServer(object):
 def test(valset: Dataset, params, model: YOLO, batch_size: int = 200) -> tuple[float, float, float, float]:
     """
     Evaluate the model on the validation dataset.
+
     Args:
         valset: validation dataset
         params: dict of parameters (must include 'names')
         model: YOLO model to evaluate
         batch_size: batch size for evaluation
+        
     Returns:
         dict with evaluation metrics (tp, fp, m_pre, m_rec, map50, mean_ap)
     """
@@ -214,7 +221,9 @@ def test(valset: Dataset, params, model: YOLO, batch_size: int = 200) -> tuple[f
     # Compute metrics
     metrics = [torch.cat(x, dim=0).cpu().numpy() for x in zip(*metrics)]  # to numpy
     if len(metrics) and metrics[0].any():
-        tp, fp, m_pre, m_rec, map50, mean_ap = util.compute_ap(*metrics, plot=False, names=params["names"])
+        tp, fp, m_pre, m_rec, map50, mean_ap = util.compute_ap(
+            *metrics, plot=False, names=params["names"]
+        )  # set plot=True to plot metric curve
     # Print results
     # print(("%10s" + "%10.3g" * 4) % ("", m_pre, m_rec, map50, mean_ap))
     # Return results