fixed total iterations for rolling training

Author: mr-yamraj

examples/pytorch/FastCells/train_classifier.py

@@ -165,24 +165,44 @@ def fit(self, training_data, validation_data, options, sparsify=False, device=No
         trim_level = options.trim_level
 
         ticks = training_data.num_rows / batch_size  # iterations per epoch
-        total_iterations = ticks * num_epochs
-        scheduler = self.configure_lr(options, optimizer, ticks, total_iterations)
-
-        # optimizer = optim.Adam(model.parameters(), lr=0.0001)
-        log = []
+        
+        # Calculation of total iterations in non-rolling vs rolling training
+        # ticks = num_rows/batch_size (total number of iterations per epoch)
+        # Non-Rolling Training:
+        # Total Iteration = num_epochs * ticks
+        # Rolling Training:
+        # irl = Initial_rolling_length (We are using 2)
+        # If num_epochs <=  max_rolling_length:
+        # Total Iterations = sum(range(irl, irl + num_epochs))
+        # If num_epochs > max_rolling_length:
+        # Total Iterations = sum(range(irl, irl + max_rolling_length)) + (num_epochs - max_rolling_length)*ticks
         if options.rolling:
             rolling_length = 2
             max_rolling_length = int(ticks)
-            if max_rolling_length > options.max_rolling_length:
-                max_rolling_length = options.max_rolling_length
+            if max_rolling_length > options.max_rolling_length + rolling_length:
+                max_rolling_length = options.max_rolling_length + rolling_length
             bag_count = 100
             hidden_bag_size = batch_size * bag_count
+            if num_epochs + rolling_length < max_rolling_length:
+                max_rolling_length = num_epochs + rolling_length
+            total_iterations = sum(range(rolling_length, max_rolling_length))
+            if num_epochs + rolling_length > max_rolling_length:
+                epochs_remaining = num_epochs + rolling_length - max_rolling_length
+                total_iterations += epochs_remaining * training_data.num_rows / batch_size
+            ticks = total_iterations / num_epochs
+        else:
+            total_iterations = ticks * num_epochs
+
+        scheduler = self.configure_lr(options, optimizer, ticks, total_iterations)
+
+        # optimizer = optim.Adam(model.parameters(), lr=0.0001)
+        log = []
 
         for epoch in range(num_epochs):
             self.train()
             if options.rolling:
                 rolling_length += 1
-                if rolling_length < max_rolling_length:
+                if rolling_length <= max_rolling_length:
                     self.init_hidden_bag(hidden_bag_size, device)
             for i_batch, (audio, labels) in enumerate(training_data.get_data_loader(batch_size)):
                 if not self.batch_first:
@@ -196,7 +216,7 @@ def fit(self, training_data, validation_data, options, sparsify=False, device=No
                 # Also, we need to clear out the hidden state,
                 # detaching it from its history on the last instance.
                 if options.rolling:
-                    if rolling_length < max_rolling_length:
+                    if rolling_length <= max_rolling_length:
                         if (i_batch + 1) % rolling_length == 0:
                             self.init_hidden()
                             break