如何将数据加载器发送到 Google Colab 中的 GPU？

青葱年少 pytorch 438

原文标题 ：How can I send a Data Loader to the GPU in Google Colab?

我有两个数据加载器，我正在尝试使用.to(device)将它们发送到 GPU，但这不起作用。这是我正在使用的代码：

# 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但这也不起作用。

原文链接：https://stackoverflow.com//questions/71527963/how-can-i-send-a-data-loader-to-the-gpu-in-google-colab

回复

我来回复
  • Joe的头像
    Joe 评论

    您不会将数据加载器移动到 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')
    ...
    2年前 0条评论

相关问题