相关环境
| 名称 | 版本 |
|---|---|
| Ubuntu | 22.04.1 |
| Python | 3.9.7 |
| PyTorch | 2.0.1+cu117 |
解决方法
问题在于utils.general中non_max_suppression函数中
1 | x[:, 5:] *= x[:, 4:5] # conf = obj_conf * cls_conf |
这一行的相乘操作创建的中间变量未能在函数执行完毕后自动释放,导致显存占用无限增长。
将其更改为:
1 | temp = x[:, 4:5].clone() |
即可及时释放显存,解决显存占用无限增长问题。