Problems with excessive GPU memory use

[jj8127]

When running test.py

for idx1, (img1_fg, img1_bg, mask_path) in enumerate(zip(loader_content_fg, loader_content_bg, mask_paths)):
    # Stylization for each source image

    print("%d/%d"%(idx1+1, len(loader_content_fg)))

    test_saver_path = os.path.join(opts.output_path, str(idx1))  
    if not os.path.exists(test_saver_path):
        os.mkdir(test_saver_path)

    # Optional. save source image
    shutil.copyfile(content_paths[idx1], os.path.join(test_saver_path, 'input.jpg'))

    # Get encoded content feature of source image        
    img1_fg = img1_fg.cuda()
    share_content_fg, _, content_fg, _ = encode_fg(img1_fg)   
    img1_bg = img1_bg.cuda()
    share_content_bg, _, content_bg, _ = encode_bg(img1_bg) 

    for idx2, (img2_fg, img2_bg) in enumerate(zip(loader_style_fg, loader_style_bg)): 
        # Stylize source image with each target style reference

        # Optional. save style image
        if idx1 == 0:
            shutil.copyfile(style_paths[idx2], os.path.join(test_saver_path, 'style_{}.jpg'.format(idx2)))

        # Get encoded style feature of target reference
        img2_fg = img2_fg.cuda()
        _, _, _, style_fg = style_encode_fg(img2_fg)
        img2_bg = img2_bg.cuda()
        _, _, _, style_bg = style_encode_bg(img2_bg)

        # Generate stylized image
        with torch.no_grad():
            outputs_fg = decode_fg(share_content_fg, content_fg, style_fg)
            outputs_bg = decode_bg(share_content_bg, content_bg, style_bg)

        # Blend images
        b,c,h,w = outputs_fg.shape #[1, 3, H, W]
        alpha = get_mask_alpha(mask_path, size=(h,w))
        outputs = blend(outputs_fg, outputs_bg, alpha)

        # Save stylized image
        path = os.path.join(test_saver_path, 'output_{}_blend.jpg'.format(idx2))
        stylized = write_image(outputs, path) # np array

        # Blending optimization
        if opts.blend_opt:
            src_img = img_as_float(imread(content_paths[idx1])) #[H, W, 3]
            mask_img = img_as_float(imread(mask_paths[idx1], as_gray=True)) # [H, W]
            stylized = stylized.cpu().detach().numpy().transpose(1, 2, 0) #[H, W, 3]
            optimized = blend_optimize(src_img, stylized, mask_img, image_size=new_size)
            # optimized = blend_optimize(src_img, stylized, mask_img, image_size=new_size, origin_res=True) #Optional. restore original high resolution
            path = os.path.join(test_saver_path, 'output_{}_opt.jpg'.format(idx2))
            imsave(path, optimized)

You cannot use more than 15GB of memory to style multiple images each time you run this iteration.

I would appreciate it if you could tell me why and how to solve it.

[chenyingshu]

Hi, may I know your image resolution during testing? We tested 512x with 11GB memory. For a larger resolution, we need more memory.