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.