我正在编写一个关于Python图像分割的youtube教程:link
本教程以其他教程为基础,我一直在参考这些教程进行优化,特别是本教程: OpenCV Pytorch Segmentation
我使用的是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都是超级内存密集型的,从我的角度来看是不可优化的。任何帮助都将不胜感激!谢谢
目前没有回答
相关问题 更多 >
编程相关推荐