Repository: OpenMined/PySyft

Issue:
TrainConfig parameter "epochs"
**TrainConfig parameter "epochs" doesn't have effect.**
After changing the number of epochs=1 to epochs=100. The worker still do only 1 epoch.

```
train_config = sy.TrainConfig(
        model=traced_model,
        loss_fn=loss_fn,
        batch_size=batch_size,
        shuffle=True,
        #max_nr_batches=max_nr_batches,
        epochs=100,
        lr=lr,
    )
```


Discussion:

Resolving patch:
diff --git a/syft/federated/federated_client.py b/syft/federated/federated_client.py
--- a/syft/federated/federated_client.py
+++ b/syft/federated/federated_client.py
@@ -72,6 +72,9 @@ def fit(self, dataset_key: str, **kwargs):
         if self.train_config is None:
             raise ValueError("TrainConfig not defined.")
 
+        if dataset_key not in self.datasets:
+            raise ValueError("Dataset {} unknown.".format(dataset_key))
+
         model = self.get_obj(self.train_config._model_id).obj
         loss_fn = self.get_obj(self.train_config._loss_fn_id).obj
 
@@ -106,19 +109,21 @@ def _fit(self, model, dataset_key, loss_fn):
 
         loss = None
         iteration_count = 0
-        for (data, target) in data_loader:
-            # Set gradients to zero
-            self.optimizer.zero_grad()
-
-            # Update model
-            output = model(data)
-            loss = loss_fn(target=target, pred=output)
-            loss.backward()
-            self.optimizer.step()
-
-            # Update and check interation count
-            iteration_count += 1
-            if iteration_count >= self.train_config.max_nr_batches >= 0:
-                break
+
+        for _ in range(self.train_config.epochs):
+            for (data, target) in data_loader:
+                # Set gradients to zero
+                self.optimizer.zero_grad()
+
+                # Update model
+                output = model(data)
+                loss = loss_fn(target=target, pred=output)
+                loss.backward()
+                self.optimizer.step()
+
+                # Update and check interation count
+                iteration_count += 1
+                if iteration_count >= self.train_config.max_nr_batches >= 0:
+                    break
 
         return loss


Test patch:
diff --git a/test/federated/test_federated_client.py b/test/federated/test_federated_client.py
--- a/test/federated/test_federated_client.py
+++ b/test/federated/test_federated_client.py
@@ -1,3 +1,5 @@
+import pytest
+
 import torch
 
 import syft as sy
@@ -55,7 +57,11 @@ def test_set_obj_other():
     assert fed_client._objects[dummy_data.id] == dummy_data
 
 
-def test_fit():
+@pytest.mark.parametrize(
+    "fit_dataset_key, epochs",
+    [("gaussian_mixture", 1), ("gaussian_mixture", 10), ("another_dataset", 1)],
+)
+def test_fit(fit_dataset_key, epochs):
     data, target = utils.create_gaussian_mixture_toy_data(nr_samples=100)
 
     fed_client = federated.FederatedClient()
@@ -100,6 +106,7 @@ def forward(self, x):
         loss_fn_id=loss_id,
         lr=0.05,
         weight_decay=0.01,
+        epochs=epochs,
     )
 
     fed_client.set_obj(model_ow)
@@ -107,32 +114,22 @@ def forward(self, x):
     fed_client.set_obj(train_config)
     fed_client.optimizer = None
 
-    for curr_round in range(12):
-        loss = fed_client.fit(dataset_key=dataset_key)
+    loss = loss_before
+    for curr_round in range(3):
+        if fit_dataset_key == dataset_key:
+            loss = fed_client.fit(dataset_key=fit_dataset_key)
+        else:
+            with pytest.raises(ValueError):
+                loss = fed_client.fit(dataset_key=fit_dataset_key)
         if PRINT_IN_UNITTESTS and curr_round % 4 == 0:  # pragma: no cover
             print("-" * 50)
             print("Iteration %s: alice's loss: %s" % (curr_round, loss))
 
-    new_model = fed_client.get_obj(model_id)
-    pred = new_model.obj(data)
-    loss_after = loss_fn(target=target, pred=pred)
-    if PRINT_IN_UNITTESTS:  # pragma: no cover:
-        print("Loss after training: {}".format(loss_after))
-
-    assert loss_after < loss_before
-
-
-def create_xor_data(nr_samples):  # pragma: no cover
-    with torch.no_grad():
-        data = torch.tensor([[0.0, 1.0], [1.0, 0.0], [1.0, 1.0], [0.0, 0.0]], requires_grad=True)
-        target = torch.tensor([1, 1, 0, 0], requires_grad=False)
-
-        data_len = int(nr_samples / 4 + 1) * 4
-        X = torch.zeros(data_len, 2, requires_grad=True)
-        Y = torch.zeros(data_len, requires_grad=False)
-
-        for i in range(int(data_len / 4)):
-            X[i * 4 : (i + 1) * 4, :] = data
-            Y[i * 4 : (i + 1) * 4] = target
+    if dataset_key == fit_dataset_key:
+        new_model = fed_client.get_obj(model_id)
+        pred = new_model.obj(data)
+        loss_after = loss_fn(target=target, pred=pred)
+        if PRINT_IN_UNITTESTS:  # pragma: no cover:
+            print("Loss after training: {}".format(loss_after))
 
-    return X, Y.long()
+        assert loss_after < loss_before