如何將數據加載器發送到 Google Colab 中的 GPU？

Moodhi 发表于 Dev

穆迪

我有兩個數據加載器，我正在嘗試將它們發送到 GPU，.to(device)但這不起作用。這是我正在使用的代碼：

# to create a batch iterator

class MyData(Dataset):
    def __init__(self, X, y):
        self.data = X
        self.target = y
        # TODO: convert this into torch code is possible
        self.length = [ np.sum(1 - np.equal(x, 0)) for x in X]
        
    def __getitem__(self, index):
        x = self.data[index]
        y = self.target[index]
        x_len = self.length[index]
        xx_len = torch.tensor(x_len)
        return {"src": x, "trg": y, "x_len": xx_len}
    
    def __len__(self):
        return len(self.data)


dataset = DataLoader(train_dataset, batch_size = BATCH_SIZE, 
                     drop_last=True,
                     shuffle=True)
test_Dataset= DataLoader(val_dataset, batch_size = BATCH_SIZE, 
                     drop_last=True,
                     shuffle=True)

我也嘗試使用pin_memory = True，但這也不起作用。

喬

您不會將數據加載器移動到 GPU。相反，創建存儲數據的批量張量，然後將這些張量移動到 GPU。

train_dataloader = DataLoader(MyData, batch_size=BS)
...
def train(nn, optim, train_dataloader, etc...): 
    for batch in train_dataloader:
        batch = batch.to('cuda')
    ...

本文收集自互联网，转载请注明来源。

如有侵权，请联系 [email protected] 删除。