test(NeuralNetwork): testing for repeatition capabilities in datagene…

…rators
frankkramer-lab · Feb 15, 2023 · c70fc4c · c70fc4c
1 parent 8f1a2e8
commit c70fc4c
Showing 1 changed file with 7 additions and 0 deletions.
diff --git a/tests/test_neuralnetwork.py b/tests/test_neuralnetwork.py
@@ -72,6 +72,13 @@ def test_training_pure(self):
                            epochs=3)
         self.assertTrue("loss" in hist)
 
+    def test_training_iterations(self):
+        model = NeuralNetwork(n_labels=4, channels=3, batch_queue_size=1)
+        hist = model.train(training_generator=self.datagen,
+                           epochs=5, iterations=10)
+        self.assertTrue("loss" in hist)
+        self.assertTrue(len(hist["loss"]) == 5)
+
     def test_training_validation(self):
         model = NeuralNetwork(n_labels=4, channels=3, batch_queue_size=1)
         hist = model.train(training_generator=self.datagen,