HongChuang/ultralytics-ascend
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions 3

Implement gc.collect() to free memory (#10129)

Browse source 
Signed-off-by: Glenn Jocher <glenn.jocher@ultralytics.com>
Co-authored-by: UltralyticsAssistant <web@ultralytics.com>
This commit is contained in:
Glenn Jocher 2024-04-17 19:37:40 -07:00 • committed by GitHub
parent 37ffebd690
commit 1f4bed233a
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 7 additions and 0 deletions
Download patch file Download diff file Expand all files Collapse all files

2
ultralytics/engine/exporter.py
View file

@ -50,6 +50,7 @@ TensorFlow.js:
$ npm start $ npm start
""" """
import gc
import json import json
import os import os
import shutil import shutil
@ -713,6 +714,7 @@ class Exporter:
# Free CUDA memory # Free CUDA memory
del self.model del self.model
gc.collect()
torch.cuda.empty_cache() torch.cuda.empty_cache()
# Write file # Write file

3
ultralytics/engine/trainer.py
View file

@ -6,6 +6,7 @@ Usage:
$ yolo mode=train model=yolov8n.pt data=coco128.yaml imgsz=640 epochs=100 batch=16 $ yolo mode=train model=yolov8n.pt data=coco128.yaml imgsz=640 epochs=100 batch=16
""" """
import gc
import math import math
import os import os
import subprocess import subprocess
@ -437,6 +438,7 @@ class BaseTrainer:
self.scheduler.last_epoch = self.epoch # do not move self.scheduler.last_epoch = self.epoch # do not move
self.stop |= epoch >= self.epochs # stop if exceeded epochs self.stop |= epoch >= self.epochs # stop if exceeded epochs
self.run_callbacks("on_fit_epoch_end") 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 torch.cuda.empty_cache() # clear GPU memory at end of epoch, may help reduce CUDA out of memory errors
# Early Stopping # Early Stopping
@ -458,6 +460,7 @@ class BaseTrainer:
if self.args.plots: if self.args.plots:
self.plot_metrics() self.plot_metrics()
self.run_callbacks("on_train_end") self.run_callbacks("on_train_end")
gc.collect()
torch.cuda.empty_cache() torch.cuda.empty_cache()
self.run_callbacks("teardown") self.run_callbacks("teardown")

2
ultralytics/utils/torch_utils.py
View file

@ -1,5 +1,6 @@
# Ultralytics YOLO 🚀, AGPL-3.0 license # Ultralytics YOLO 🚀, AGPL-3.0 license
import gc
import math import math
import os import os
import random import random
@ -581,6 +582,7 @@ def profile(input, ops, n=10, device=None):
except Exception as e: except Exception as e:
LOGGER.info(e) LOGGER.info(e)
results.append(None) results.append(None)
gc.collect() # attempt to free unused memory
torch.cuda.empty_cache() torch.cuda.empty_cache()
return results return results