我想在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_x
和data_y
行更改为:
data_x = torch.randn(batch_size, n_input).cuda()
data_y = (torch.rand(size=(batch_size, 1)) < 0.5).float().cuda()
但是我得到了同样的错误-有人能解释一下如何在GPU上运行这个代码吗?
1条答案
按热度按时间irlmq6kh1#
正如评论中所建议的,您需要将模型和数据传输到同一个设备上。