Pikaliov/caption-test
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add VRAM cleanup (gc + empty_cache) before training start

Browse Source 
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
pikaliov
2026-04-21 19:36:25 +03:00
parent 858718431b
commit e3ecb09687
1 changed files with 12 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
src/training/train_gtauav.py
View File

@@ -200,6 +200,17 @@ def _evaluate(
return metrics return metrics
def _clear_vram() -> None:
"""Free VRAM from previous runs before starting."""
import gc
gc.collect()
if torch.cuda.is_available():
torch.cuda.empty_cache()
torch.cuda.reset_peak_memory_stats()
allocated = torch.cuda.memory_allocated() / 1e9
LOGGER.info("🧹 VRAM cleared. Current usage: %.2f GB", allocated)
def train(cfg: TrainConfigGTAUAV) -> None: def train(cfg: TrainConfigGTAUAV) -> None:
"""Run full training loop.""" """Run full training loop."""
coloredlogs.install( coloredlogs.install(
@@ -207,6 +218,7 @@ def train(cfg: TrainConfigGTAUAV) -> None:
logger=LOGGER, logger=LOGGER,
fmt="%(asctime)s %(name)s %(levelname)s %(message)s", fmt="%(asctime)s %(name)s %(levelname)s %(message)s",
) )
_clear_vram()
_set_seed(cfg.seed) _set_seed(cfg.seed)
output_dir = Path(cfg.output_dir) output_dir = Path(cfg.output_dir)
output_dir.mkdir(parents=True, exist_ok=True) output_dir.mkdir(parents=True, exist_ok=True)