mirrors
/
PaddleClas
mirror of https://github.com/PaddlePaddle/PaddleClas.git
1
0
Fork 0
Code Issues Projects Releases Wiki Activity

fix gpu memory growth

pull/3281/head
zhangyubo0722 2024-10-18 04:22:06 +00:00 committed by Tingquan Gao
parent 9f0cf7a7eb
commit d178931fca
1 changed files with 7 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
ppcls/engine/engine.py
View File

@ -15,6 +15,7 @@ from __future__ import division
from __future__ import print_function
import os
import gc
import shutil
import copy
import platform
@ -407,10 +408,12 @@ class Engine(object):
save_path = os.path.join(self.output_dir, prefix,
"inference")
self.export(save_path, uniform_output_enabled)
gc.collect()
if self.ema:
ema_save_path = os.path.join(
self.output_dir, prefix, "inference_ema")
self.export(ema_save_path, uniform_output_enabled)
gc.collect()
update_train_results(
self.config, prefix, metric_info, ema=self.ema)
save_load.save_model_info(metric_info, self.output_dir,
@ -436,10 +439,12 @@ class Engine(object):
save_path = os.path.join(self.output_dir, prefix,
"inference")
self.export(save_path, uniform_output_enabled)
gc.collect()
if self.ema:
ema_save_path = os.path.join(self.output_dir, prefix,
"inference_ema")
self.export(ema_save_path, uniform_output_enabled)
gc.collect()
update_train_results(
self.config,
prefix,
@ -464,10 +469,12 @@ class Engine(object):
if uniform_output_enabled:
save_path = os.path.join(self.output_dir, prefix, "inference")
self.export(save_path, uniform_output_enabled)
gc.collect()
if self.ema:
ema_save_path = os.path.join(self.output_dir, prefix,
"inference_ema")
self.export(ema_save_path, uniform_output_enabled)
gc.collect()
save_load.save_model_info(metric_info, self.output_dir, prefix)
self.model.train()