1. 程式人生 > 其它 >【深度學習】RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!

【深度學習】RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!

報錯程式碼:

if __name__ == '__main__':
    model = Perception(2, 3, 2).cuda()

    input = torch.randn(4, 2).cuda()
    output = model(input)
    # output = output.cuda()

    label = torch.Tensor([0, 1, 1, 0]).long()
    criterion = nn.CrossEntropyLoss()
    loss_nn = criterion(output, label)
    print(loss_nn)
    loss_functional = F.cross_entropy(output, label)
    print(loss_functional)

報錯截圖如下:

RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu! (when checking argument for argument target in method wrapper_nll_loss_forward)

報這個錯的原因在於,程式碼中的Tensor,一會在CPU中執行,一會在GPU中執行,所以最好是都放在同一個device中執行

核心程式碼:

device = torch.device('cuda:0')

並且將用到的Tensor都改為同一個device:Tensor.to(device)

上述程式碼修改後:

if __name__ == '__main__':
    device = torch.device('cuda:0')
    model = Perception(2, 3, 2).to(device)

    input = torch.randn(4, 2).to(device)
    output = model(input).to(device)

    label = torch.Tensor([0, 1, 1, 0]).long().to(device)
    criterion = nn.CrossEntropyLoss()
    loss_nn = criterion(output, label).to(device)
    print(loss_nn)
    loss_functional = F.cross_entropy(output, label)
    print(loss_functional)

這樣就不會報錯了

完整程式碼:

import torch
from torch import nn
import torch.nn.functional as F

from torch.nn import Linear


class linear(nn.Module):  # 繼承nn.Module
    def __init__(self, in_dim, out_dim):
        super(Linear, self).__init__()  # 呼叫nn.Module的建構函式

        # 使用nn.Parameter來構造需要學習的引數
        self.w = nn.Parameter(torch.randn(in_dim, out_dim))
        self.b = nn.Parameter(torch.randn(out_dim))

    # 在forward中實現前向傳播過程
    def forward(self, x):
        x = x.matmul(self.w)
        y = x + self.b.expand_as(x)  # expand_as保證矩陣形狀一致

        return y

class Perception(nn.Module):
    def __init__(self, in_dim, hid_dim, out_dim):
        super(Perception, self).__init__()
        self.layer = nn.Sequential(
            nn.Linear(in_dim, hid_dim),
            nn.Sigmoid(),
            nn.Linear(hid_dim, out_dim),
            nn.Sigmoid()
        )
        # self.layer1 = Linear(in_dim, hid_dim)
        # self.layer2 = Linear(hid_dim, out_dim)

    def forward(self, x):
        # x = self.layer1(x)
        # y = torch.sigmoid(x)
        # y = self.layer2(y)
        # y = torch.sigmoid(y)
        y = self.layer(x)
        return y


if __name__ == '__main__':
    device = torch.device('cuda:0')
    model = Perception(2, 3, 2).to(device)

    input = torch.randn(4, 2).to(device)
    output = model(input).to(device)
    # output = output.cuda()

    label = torch.Tensor([0, 1, 1, 0]).long().to(device)
    criterion = nn.CrossEntropyLoss()
    loss_nn = criterion(output, label).to(device)
    print(loss_nn)
    loss_functional = F.cross_entropy(output, label)
    print(loss_functional)