From 1f4bed233a3cab515f32fe30f667c746f77a3ff2 Mon Sep 17 00:00:00 2001
From: Glenn Jocher <glenn.jocher@ultralytics.com>
Date: Wed, 17 Apr 2024 19:37:40 -0700
Subject: [PATCH] Implement `gc.collect()` to free memory (#10129)

Signed-off-by: Glenn Jocher <glenn.jocher@ultralytics.com>
Co-authored-by: UltralyticsAssistant <web@ultralytics.com>
---
 ultralytics/engine/exporter.py   | 2 ++
 ultralytics/engine/trainer.py    | 3 +++
 ultralytics/utils/torch_utils.py | 2 ++
 3 files changed, 7 insertions(+)

diff --git a/ultralytics/engine/exporter.py b/ultralytics/engine/exporter.py
index d38738433a..0beb38b960 100644
--- a/ultralytics/engine/exporter.py
+++ b/ultralytics/engine/exporter.py
@@ -50,6 +50,7 @@ TensorFlow.js:
     $ npm start
 """
 
+import gc
 import json
 import os
 import shutil
@@ -713,6 +714,7 @@ class Exporter:
 
         # Free CUDA memory
         del self.model
+        gc.collect()
         torch.cuda.empty_cache()
 
         # Write file
diff --git a/ultralytics/engine/trainer.py b/ultralytics/engine/trainer.py
index 1af3721eac..a33aac6ec1 100644
--- a/ultralytics/engine/trainer.py
+++ b/ultralytics/engine/trainer.py
@@ -6,6 +6,7 @@ Usage:
     $ yolo mode=train model=yolov8n.pt data=coco128.yaml imgsz=640 epochs=100 batch=16
 """
 
+import gc
 import math
 import os
 import subprocess
@@ -437,6 +438,7 @@ class BaseTrainer:
                 self.scheduler.last_epoch = self.epoch  # do not move
                 self.stop |= epoch >= self.epochs  # stop if exceeded epochs
             self.run_callbacks("on_fit_epoch_end")
+            gc.collect()
             torch.cuda.empty_cache()  # clear GPU memory at end of epoch, may help reduce CUDA out of memory errors
 
             # Early Stopping
@@ -458,6 +460,7 @@ class BaseTrainer:
             if self.args.plots:
                 self.plot_metrics()
             self.run_callbacks("on_train_end")
+        gc.collect()
         torch.cuda.empty_cache()
         self.run_callbacks("teardown")
 
diff --git a/ultralytics/utils/torch_utils.py b/ultralytics/utils/torch_utils.py
index 2b0ad9513b..64bd6d6f13 100644
--- a/ultralytics/utils/torch_utils.py
+++ b/ultralytics/utils/torch_utils.py
@@ -1,5 +1,6 @@
 # Ultralytics YOLO 🚀, AGPL-3.0 license
 
+import gc
 import math
 import os
 import random
@@ -581,6 +582,7 @@ def profile(input, ops, n=10, device=None):
             except Exception as e:
                 LOGGER.info(e)
                 results.append(None)
+            gc.collect()  # attempt to free unused memory
             torch.cuda.empty_cache()
     return results