PyTorch在Apple Silicon上DataLoader的CPU与MPS结果不一致问题

问题描述 #

在Apple Silicon Mac上使用PyTorch的DataLoader时，切换到MPS后端后出现标签解析异常：CPU后端下张量y能被正确读取，但MPS后端返回的batch_labels长度正常，所有元素却与y的首个元素值完全相同。

复现代码 #

import torch
from torch.utils.data import TensorDataset, random_split, DataLoader


device = torch.device("mps") 
X = torch.tensor([[[0.5,0.4], [0,0]],[[0.3,0.2], [0,0]],[[0.5,0.2], [0,0]],[[0.2,0.2], [0,0]]], dtype=torch.float32).to(device)
y = torch.tensor([1,0,0,0], dtype=torch.float32).to(device)

print(X.shape)
print(y.shape)
print(y)
dataset = TensorDataset(X, y)
train_size = int(0.5 * len(dataset))
test_size = len(dataset) - train_size
train_dataset, test_dataset = random_split(dataset, [train_size, test_size])
train_loader = DataLoader(train_dataset, batch_size=10, shuffle=True)

for i, (batch_data, batch_labels) in enumerate(train_loader):
    print(batch_data)
    print(batch_labels)
    break

MPS后端下的异常输出 #

torch.Size([4, 2, 2])
torch.Size([4])
tensor([1., 0., 1., 0.], device='mps:0')
tensor([[[0.5000, 0.2000],
         [0.0000, 0.0000]],

        [[0.5000, 0.4000],
         [0.0000, 0.0000]]], device='mps:0')
tensor([1., 1.], device='mps:0')

问题说明 #

该问题是PyTorch针对MPS后端实现的已知bug，与MPS算子覆盖跟踪逻辑相关。

内容的提问来源于stack exchange，提问作者Pavol Travnik