models: Fix runtime tracking in multi-sample prediction.

fukuda-lab · Jul 13, 2023 · 3061eac · 3061eac
1 parent 988bfd0
commit 3061eac
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 0 deletions.
diff --git a/code/models/MLPAutoEncoder.py b/code/models/MLPAutoEncoder.py
@@ -112,8 +112,10 @@ def predict(self, data: EncodedSampleGenerator, **kwargs) -> SampleGenerator:
             prediction = self.model_instance.predict(encoded_sample)
 
             if isinstance(sample, list):
+                sum_processing_time += time.process_time_ns() - start_time_ref
                 # Handle the prediction for multi-sample encoding.
                 for i, sample in enumerate(sample):
+                    start_time_ref = time.process_time_ns()
                     sample[PredictionField.MODEL_NAME] = self.model_name
                     sample[PredictionField.OUTPUT_DISTANCE] = prediction[i]
                     sum_processing_time += time.process_time_ns() - start_time_ref

diff --git a/code/models/RandomForest.py b/code/models/RandomForest.py
@@ -94,13 +94,18 @@ def predict(self, data: EncodedSampleGenerator, **kwargs) ->SampleGenerator:
         for sample, encoded_sample in data:
             start_time_ref = time.process_time_ns()
             prediction = self.model_instance.predict(encoded_sample)
+
             if isinstance(sample, list):
+                sum_processing_time += time.process_time_ns() - start_time_ref
+                # Handle the prediction for multi-sample encoding.
                 for i, sample in enumerate(sample):
+                    start_time_ref = time.process_time_ns()
                     sample[PredictionField.MODEL_NAME] = self.model_name
                     sample[PredictionField.OUTPUT_BINARY] = prediction[i]
                     sum_processing_time += time.process_time_ns() - start_time_ref
                     sum_samples += 1
                     yield sample
+
             else:
                 sample[PredictionField.MODEL_NAME] = self.model_name
                 sample[PredictionField.OUTPUT_BINARY] = prediction[0]