Speeding up a script that uses Img2Vec library

vision
1

Is there a way to run this across 4 GPUs or speed it up any other way? The script takes a really long time. Also, I have 4 GPUs each with 11G of GPU memory. However, this code only takes ~1G of one of the GPUs.

from img2vec_pytorch import Img2Vec
from PIL import Image
import torch
import glob
import os

# # Initialize Img2Vec with GPU
img2vec = Img2Vec(cuda=True)
for dirpath, dirname, filename in os.walk('/SeaExp/ABC/'):
    if dirpath.endswith('XYZ'):
        images_full_path = dirpath + '/*.jpeg'
        for img in glob.glob(images_full_path):
            rgb_img = Image.open(img)
            vec = img2vec.get_vec(rgb_img, tensor=True)
            vec_name = os.path.basename(img)[:-5] + '.pt'
            dir_name = os.path.dirname(img) 
            torch.save(vec, dir_name + '/' + vec_name)

cross-posted here since I thought Iā€™d get more feedback from PyTorch community.

2

Try to pass the images as a list instead of each image separately.
Batching the workload would generally increase the device utilization and would yield a speedup.

1 Like