CODEWITHSUNDEEP
pytorch
Slowat_Kela
Slowat_Kela

Reputation: 1511

How to convert pytorch model to being on GPU?

I want to run pytorch on a GPU.

I have this code:

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()

When I run it, the output is:

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

when I uncomment the line model.to(device) and rerun it, I get:

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)

I can see other questions like this (e.g. here) but can't work out how I'm meant to convert to GPU.

For example, I change the data_x and data_y lines to:

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

But I get the same error - could someone explain how to run this code on a GPU?

Upvotes: 1

Views: 1864

Answers (1)

Matt
Matt

Reputation: 4049

As suggested in the comments, you need to transfer both your model and your data to the same device. Below should work:

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()

Upvotes: 2

Related Questions