如何将pytorch模型转换到GPU上?

ogsagwnx  于 2022-11-09  发布在  其他
关注(0)|答案(1)|浏览(214)

我想在GPU上运行pytorch。
我有这样的代码:

import torch
import torch.nn as nn
device = torch.device("cuda:0")

n_input, n_hidden, n_out, batch_size, learning_rate = 10, 15, 1, 100, 0.01

data_x = torch.randn(batch_size, n_input)
data_y = (torch.rand(size=(batch_size, 1)) < 0.5).float()

print(data_x.size())
print(data_y.size())

model = nn.Sequential(nn.Linear(n_input, n_hidden),
                      nn.ReLU(),
                      nn.Linear(n_hidden, n_out),
                      nn.Sigmoid())

# model.to(device)

print(next(model.parameters()).is_cuda)

loss_function = nn.MSELoss()
optimizer = torch.optim.SGD(model.parameters(), lr=learning_rate)

losses = []
for epoch in range(5000):
    pred_y = model(data_x)
    loss = loss_function(pred_y, data_y)
    losses.append(loss.item())

    model.zero_grad()
    loss.backward()
    optimizer.step()

print(torch.cuda.get_device_name())
print(torch.__version__)
print(torch.version.cuda)

import matplotlib.pyplot as plt
plt.plot(losses)
plt.ylabel('loss')
plt.xlabel('epoch')
plt.title("Learning rate %f"%(learning_rate))
plt.show()

当我运行它时,输出是:

torch.Size([100, 10])
torch.Size([100, 1])
False
Quadro P2000
1.12.1
11.3

当我取消model.to(device)行的注解并重新运行它时,我得到:

Traceback (most recent call last):
  File "basic_pytorch_with_gpu.py", line 26, in <module>
    pred_y = model(data_x)
  File "/opt/conda/lib/python3.7/site-packages/torch/nn/modules/module.py", line 1130, in _call_impl
    return forward_call(*input,**kwargs)
  File "/opt/conda/lib/python3.7/site-packages/torch/nn/modules/container.py", line 139, in forward
    input = module(input)
  File "/opt/conda/lib/python3.7/site-packages/torch/nn/modules/module.py", line 1130, in _call_impl
    return forward_call(*input,**kwargs)
  File "/opt/conda/lib/python3.7/site-packages/torch/nn/modules/linear.py", line 114, in forward
    return F.linear(input, self.weight, self.bias)
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 mat1 in method wrapper_addmm)

我可以看到其他类似的问题(例如here),但无法计算出如何转换到GPU。
例如,我将data_xdata_y行更改为:

data_x = torch.randn(batch_size, n_input).cuda()
data_y = (torch.rand(size=(batch_size, 1)) < 0.5).float().cuda()

但是我得到了同样的错误-有人能解释一下如何在GPU上运行这个代码吗?

irlmq6kh

irlmq6kh1#

正如评论中所建议的,您需要将模型和数据传输到同一个设备上。

import torch
import torch.nn as nn
device = torch.device("cuda:0")

n_input, n_hidden, n_out, batch_size, learning_rate = 10, 15, 1, 100, 0.01

data_x = torch.randn(batch_size, n_input)
data_y = (torch.rand(size=(batch_size, 1)) < 0.5).float()

print(data_x.size())
print(data_y.size())

model = nn.Sequential(nn.Linear(n_input, n_hidden),
                      nn.ReLU(),
                      nn.Linear(n_hidden, n_out),
                      nn.Sigmoid())

# Transfer to device

model = model.to(device)
data_x = data_x.to(device)
data_y = data_y.to(device)

print(next(model.parameters()).is_cuda)

loss_function = nn.MSELoss()
optimizer = torch.optim.SGD(model.parameters(), lr=learning_rate)

losses = []
for epoch in range(5000):
    pred_y = model(data_x)
    loss = loss_function(pred_y, data_y)
    losses.append(loss.item())

    model.zero_grad()
    loss.backward()
    optimizer.step()

print(torch.cuda.get_device_name())
print(torch.__version__)
print(torch.version.cuda)

import matplotlib.pyplot as plt
plt.plot(losses)
plt.ylabel('loss')
plt.xlabel('epoch')
plt.title("Learning rate %f"%(learning_rate))
plt.show()

相关问题