Dataloader使用mps - PyTorch的速度很慢。

Question

由于某些原因，当使用mps时，数据存储程序要慢得多(在这一点上使用cpu更好)。

知道为什么吗？

复制代码：

class Dataset(torch.utils.data.Dataset):
    def __init__(self, device):
        self.a = torch.tensor(1, device=device)
        
    def __len__(self):
        return 100
    
    def __getitem__(self, i):
        return self.a, self.a
    
for device in ['mps', 'cpu']:
    dataloader = torch.utils.data.DataLoader(Dataset(device), 64)
    %time next(iter(dataloader))

谢谢!

Answer 1

根据MPS的设置方式(已经有一段时间没有做任何金属编码了，所以.)，您的数据可能必须先通过cpu从磁盘加载到gpu。请注意，这只是数据加载时间，如果您正在进行培训和/或推断，这可能对mps更快。