CUDA OOM in test phase

vision
1

My simple test code to test the net,but it fails because the CUDA OOM error.

def test(model,test_dir):
    model.cuda()
    model.eval()
    if not os.path.exists(os.path.join(test_dir,'result')):
        os.makedirs(os.path.join(test_dir,'result'))
    for file in os.listdir(test_dir):
        with torch.no_grad():
            x=Image.open(os.path.join(test_dir,file))
            x.resize((640,480),Image.BILINEAR)
            x=to_tensor(x)
            x.unsqueeze_(0)
            x=x.cuda()
            y=model(x)
            y=torch.argmax(F.softmax(y,1),1)
            y=y.cpu().float()
            y=to_pil_image(y)
            y.save(os.path.join(test_dir,'result',file.strip('jpg')+'png'))
            del x,y

model=Model()
load_checkpoint(epoch=450,[('model', model)])
seg=model.segmentation.state_dict()
model=Segmentaion()
model.load_weights(seg)
test(model,'./test/origin')
2

I use nvidia-smi to supervise the mem usage of gpu card.
My 1080ti is 3373mb used in the first nearly 20 pics,then turns to about 8Gb occupied ,and then OOM,I dont know why

3

I can’t see any critical problem with this code, but some small changes you could try:

  1. move the with torch.no_grad() out of for loop
  2. if the oom happens exactly after 20 pics, you could set a breakpoint at that moment and have a closer look at the variables.
  3. maybe upgrade PyTorch?