[BUG] Model Training Crashing Windows

[themetzmeier]

🧠 Describe the Bug

Possible VRAM Memory Leak in detector model training causes Windows to crash. Training of a large model crashed with only 60 total images and a batch size of 2 on a NVIDIA 3080Ti with 12GB of VRAM. This doesn't happen consistently either, as this is the first time this has happened and I have ran these regression tests weekly for the past few weeks without issue on the same dataset, same kwargs.

🔁 Steps to Reproduce

lightly_train.train_object_detection(**train_kwargs)

🤖 Environment Details

• OS: Windows 11
• Python version: Python3.12
• Frameworks/Libraries (with versions):
  torch 2.11.0+cu128
  lightly_train 0.14.3
  cuda 12.8

📌 Additional Context

Output from WinDbg:
VIDEO_TDR_FAILURE (0x116)
nvlddmkm.sys (NVIDIA driver)
0xc000009a (insufficient resources)

[yutong-xiang-97]

Hi @themetzmeier, thank you for raising the issue!

This memory leak most likely comes from the variable image sizes from the scale jitter augmentation. The largest size used in scale jitter can be (812, 812) (see here). With a large model, one can get a memory peak with this size applied that cause VRAM leakage.

If it is a routine test, I would recommend you pass constant scale jitter sizes that works for you, e.g. transform_args={"scale_jitter": {"sizes": [(640,640)]}}.

Hope this helps!