Python 如何使用Resnet、FCn、DeepLab图像分割评估解决GPU OOM问题？

我正在编写一个关于Python图像分割的youtube教程：

本教程以其他教程为基础，我一直在参考这些教程进行优化，特别是本教程：

我使用的是NVDIA2070图形卡，带有8GB的GPU内存

我的问题是，最初的教程通过FCN教授了使用Resnet的语义分段程序的基本CPU实现。我想在它的基础上利用GPU，所以我找到了后面的教程。我在这方面没有任何经验，但我知道如何在GPU上运行它，并立即遇到GPU OOM问题：

---

运行时错误：CUDA内存不足。试图分配184.00 MiB （GPU 0；8.00 GiB总容量；5.85 GiB已分配；26.97 MiB空闲；PyTorch总共保留5.88 GiB）

---

当我在一个小图像上运行这个程序，或者我将高清图像的图像质量降低到50%分辨率时，我并没有得到OOM错误

我的戳戳让我相信，我的OOM是如何在整个任务中分配内存的结果。所以现在我尝试实现了另一种DeepLab解决方案，希望它能更有效地分配内存，但事实并非如此

这是我的密码：

from PIL import Image
import torch
import torchvision.transforms as T
from torchvision import models
import numpy as np
import imghdr

fcn = None
dlab = None

def getRotoModel():
    global fcn
    global dlab
    fcn = models.segmentation.fcn_resnet101(pretrained=True).eval()
    dlab = models.segmentation.deeplabv3_resnet101(pretrained=1).eval()

# Define the helper function
def decode_segmap(image, nc=21):

    label_colors = np.array([(0, 0, 0),  # 0=background
                           # 1=aeroplane, 2=bicycle, 3=bird, 4=boat, 5=bottle
               (128, 0, 0), (0, 128, 0), (128, 128, 0), (0, 0, 128), (128, 0, 128),
               # 6=bus, 7=car, 8=cat, 9=chair, 10=cow
               (0, 128, 128), (128, 128, 128), (64, 0, 0), (192, 0, 0), (64, 128, 0),
               # 11=dining table, 12=dog, 13=horse, 14=motorbike, 15=person
               (192, 128, 0), (64, 0, 128), (192, 0, 128), (64, 128, 128), (192, 128, 128),
               # 16=potted plant, 17=sheep, 18=sofa, 19=train, 20=tv/monitor
               (0, 64, 0), (128, 64, 0), (0, 192, 0), (128, 192, 0), (0, 64, 128)])

    r = np.zeros_like(image).astype(np.uint8)
    g = np.zeros_like(image).astype(np.uint8)
    b = np.zeros_like(image).astype(np.uint8)

    for l in range(0, nc):
        idx = image == l
        r[idx] = label_colors[l, 0]
        g[idx] = label_colors[l, 1]
        b[idx] = label_colors[l, 2]

    rgb = np.stack([r, g, b], axis=2)
    return rgb

valid_images = ['jpg','png', 'rgb', 'pbm', 'ppm', 'tiff', 'rast', 'xbm', 'bmp', 'exr', 'jpeg'] #Valid image formats
dev = torch.device('cuda')
def createMatte(filename, matteName, factor):
    if imghdr.what(filename) in valid_images:
        img = Image.open(filename).convert('RGB')
        
        size = img.size
        w, h = size
        modifiedSize = h * factor
        print('Image original size is ', size)
        print('Modified size is ', modifiedSize)
        trf = T.Compose([T.Resize(int(modifiedSize)),
                     T.ToTensor(), 
                     T.Normalize(mean = [0.485, 0.456, 0.406], 
                                 std = [0.229, 0.224, 0.225])])
        inp = trf(img).unsqueeze(0)
        #inp = trf(img).unsqueeze(0).to(dev)
        
        if (fcn == None): getRotoModel()
        
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
            inp = inp.to(dev)
            fcn.to(dev)
            out = fcn.to(dev)(inp)['out'][0]
        
        with torch.no_grad():
            out = fcn(inp)['out'][0]
        
        #out = fcn(inp)['out']
        #out = fcn.to(dev)(inp)['out']
        om = torch.argmax(out.squeeze(), dim=0).detach().cpu().numpy()  
        rgb = decode_segmap(om)
        im = Image.fromarray(rgb)
        im.save(matteName)
    else:
        print('File type is not supported for file ' + filename)
        print(imghdr.what(filename))
        
def createDLMatte(filename, matteName, factor):
    if imghdr.what(filename) in valid_images:
        img = Image.open(filename).convert('RGB')
            
        size = img.size
        w, h = size
        modifiedSize = h * factor
        print('Image original size is ', size)
        print('Modified size is ', modifiedSize)
        trf = T.Compose([T.Resize(int(modifiedSize)),
            T.ToTensor(), 
            T.Normalize(mean = [0.485, 0.456, 0.406], 
                        std = [0.229, 0.224, 0.225])])
        inp = trf(img).unsqueeze(0)
        #inp = trf(img).unsqueeze(0).to(dev)
            
        if (dlab == None): getRotoModel()
            
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
            inp = inp.to(dev)
            dlab.to(dev)
            out = dlab.to(dev)(inp)['out'][0]
            
        with torch.no_grad():
            out = dlab(inp)['out'][0]
            
        #out = fcn(inp)['out']
        #out = fcn.to(dev)(inp)['out']
        om = torch.argmax(out.squeeze(), dim=0).detach().cpu().numpy()  
        rgb = decode_segmap(om)
        im = Image.fromarray(rgb)
        im.save(matteName)        

我想知道的是，GPU问题有解决方案吗？我不想把自己局限于CPU渲染，当我有一个功能强大的GPU时，每幅图像大约需要一分钟。正如我所说，我对这方面的大部分内容都很陌生，但我希望有一种方法可以在这个过程中更好地分配内存

我有一些潜在的解决方案，但我无法找到用于实施的资源

（糟糕的解决方案）当GPU接近内存结束时，在GPU上设置计算上限，并将剩余任务切换到CPU。我不仅觉得这是一个糟糕的问题，我也不知道如何在任务中实现GPU CPU切换

（更好）通过将图像分割成可管理的位，并将这些位保存到临时文件中，然后最终将它们合并，来修复内存分配

两者的某种结合

现在我担心分割图像会降低结果的质量，因为每一块图像都不符合上下文，我需要某种智能拼接，而这超出了我的工资等级

所以我通常会问，是否有资源来解决这些可能的解决方案，或者是否有更好的解决方案

最后，我的实现是否存在导致GPU OOM错误的错误？我不知道是我的代码没有优化，还是DeepLab和FCN都是超级内存密集型的，从我的角度来看是不可优化的。任何帮助都将不胜感激！谢谢