训练模型时候显存爆炸的一种可能性以及解决办法

转载：https://wstchhwp.blog.csdn.net/article/details/108405102
模型训练的时候，爆显存了，可以调整batch，对数据进行crop等等操作。
今天发现一个模型，训练ok，每次测试的时候爆显存。开始以为是因为用了全图（1920x1080略大）进行inference，这是一方面。但后来发现忘了用with torch.no_grad():这导致模型运算的时候不能释放显存（记录了梯度信息），所以显存巨大。加了之后，用了不过3G显存就够了。确实inference不需要那么多显存的，以后记着这种不正常现象如何处理。

一般训练不爆显存，测试也不会爆；训练时的显存占用远多于inference

with torch.no_grad():
                output = self.model(lr)`