[Python] Loading many PyTorch .pt files from Google Drive in Google Colab is extremely slow

I am storing my ESM protein embeddings as individual .pt files in Google Drive (about 46,000 files, total size ~182 MB).

When I try to load them in Google Colab, the process is extremely slow (only ~2 files per second).

Here is my code:

import glob
import torch
from tqdm import tqdm

folder = "drive/MyDrive/train_embeddings_35M/*.pt"

id_to_emb = {}
files = glob.glob(folder)

print("Total files:", len(files))

for file in tqdm(files, desc="Loading embeddings"):
data = torch.load(file, map_location="cuda")

raw_id = data["entry_id"]
formatted_id = raw_id.split("|")[1]

layer = list(data["mean_representations"].keys())[0]
emb = data["mean_representations"][layer]

id_to_emb[formatted_id] = emb


That's my output:

Total files: 46181
Loading embeddings: 3%|▎ | 1380/46181 [07:50<5:00:07, 2.49it/s]


Is there any better way to store them or loading them?

Continue reading...