def calculate_similarity(model, image_1, image_2, draw=False):
    image_1_pil = Image.fromarray(image_1)
    image_2_pil = Image.fromarray(image_2)
    
    probability = model.detect_image(image_1_pil, image_2_pil, draw)
    '''
    每次跑完显存代码都释放显存
    '''
    torch.cuda.empty_cache()
    return probability

        当我写了一个多线程代码,但是没有释放显存的习惯时:

占显存类代码随手丢垃圾习惯的重要性-LMLPHP

疯狂的爆OOM,

RuntimeError: CUDA out of memory. Tried to allocate 530.00 MiB (GPU 0; 2.00 GiB total capacity; 605.78 MiB already allocated; 206.31 MiB free; 612.00 MiB reserved in total by PyTorch)

非常的痛苦,

那么我们如何释放显存呢?

1. 把需要使用显存的代码,也就是需要调用CUDA的代码,打包为一个函数。

2.在函数的结尾加上。

'''
torch.cuda.empty_cache()`是PyTorch提供的函数,
用于清空当前GPU上的缓存。在深度学习模型训练过程中,
GPU显存可能会被一些中间结果占用而没有及时释放,
导致显存不足的问题。调用torch.cuda.empty_cache()可以手动清理这些未释放的显存,
确保GPU显存能够被有效地管理和利用。
这个函数通常在模型训练的迭代过程中的适当时机调用,
以避免显存占用过高而导致程序崩溃。
'''
torch.cuda.empty_cache()

  torch.cuda.empty_cache()`是PyTorch提供的函数,用于清空当前GPU上的缓存。在深度学习模型训练过程中,GPU显存可能会被一些中间结果占用而没有及时释放,导致显存不足的问题。调用torch.cuda.empty_cache()可以手动清理这些未释放的显存,确保GPU显存能够被有效地管理和利用。这个函数通常在模型训练的迭代过程中的适当时机调用,以避免显存占用过高而导致程序崩溃。

04-16 00:21