Memory usage during pre-processing

[koustav123]

Hi,
It takes more than 100GBs of RAM during pre-processing the PNGs. Is there a less memory intensive version available somewhere ?
Thanks in advance

[glample]

Hi,
It should not require that much, I'll have a look, but I agree that the way it is done currently is not optimal.. Are you working with 256 resolution images? The problem is that 200k images in RGB, of size 256x256, will necessarily take ~40GB in memory... With 128 images you would "only" need 10 GB RAM to store the entire dataset in memory. An alternative would be to load images on the fly, but this will be slower.

[LynnHo]

@koustav123 replacing the original function by this one reduces memory

def preprocess_images():

    if os.path.isfile(IMG_PATH):
        print("%s exists, nothing to do." % IMG_PATH)
        return

    print("Reading images from img_align_celeba/ ...")
    data = np.zeros([N_IMAGES, 3, IMG_SIZE, IMG_SIZE], dtype=np.uint8)
    for i in range(1, N_IMAGES + 1):
        if i % 10000 == 0:
            print(i)
        image = mpimg.imread('img_align_celeba/%06i.jpg' % i)[20:-20]
        assert image.shape == (178, 178, 3)
        if IMG_SIZE < 178:
            image = cv2.resize(image, (IMG_SIZE, IMG_SIZE), interpolation=cv2.INTER_AREA)
        elif IMG_SIZE > 178:
            image = cv2.resize(image, (IMG_SIZE, IMG_SIZE), interpolation=cv2.INTER_LANCZOS4)
        assert image.shape == (IMG_SIZE, IMG_SIZE, 3)
        data[i - 1, ...] = image.transpose((2, 0, 1))

    data = torch.from_numpy(data)
    assert data.size() == (N_IMAGES, 3, IMG_SIZE, IMG_SIZE)

    print("Saving images to %s ..." % IMG_PATH)
    torch.save(data[:20000].clone(), 'images_%i_%i_20000.pth' % (IMG_SIZE, IMG_SIZE))
    torch.save(data, IMG_PATH)