pytorch数据加载遇到的坑

问题

pytorch 版本1.6

最近在处理scRNA-seq数据，发现数据读取贼慢，dataset，__getitem__函数如下：

    def __getitem__(self, index):
        return {"x":self.adata.X[index],
                "count":self.adata.layers["count"][index],
                "label":self.adata.obs["label"][index]}

from torch.utils.data import DataLoader
...
...
train_loader = DataLoader(dataset, shuffle=True, batch_size=128, pin_memory=True, num_workers=0)

for epoch in tqdm(range(400)):
    for batch in train_loader:
        batch = {key:value.cuda() for key, value in batch.items()}

400 epoch 耗时 20s：

100%|█████████████████████████████████████████| 400/400 [00:20<00:00, 19.75it/s]

 而scvi模型数据读取贼快，耗时2s：

from scvi.core.dat