Когда выполняется функция forward
моей нейронной сети (после завершения фазы обучения), я испытываю RuntimeError: Expected object of backend CUDA but got backend CPU for argument #4 'mat1'.
Трассировка ошибки указывает, что ошибка происходит из-за вызова команды output = self.layer1(x)
.Я попытался перенести все данные тензоров на мой графический процессор.Кажется, я тоже что-то пропускаю.
Вот код, который я пробовал:
use_cuda = torch.cuda.is_available()
device = torch.device('cuda:0' if use_cuda else 'cpu')
class NeuralNet(nn.Module):
def __init__(self, input_size, hidden_size, output_size):
super(NeuralNet, self).__init__()
self.layer1 = nn.Linear(input_size, hidden_size).cuda(device)
self.layer2 = nn.Linear(hidden_size, output_size).cuda(device)
self.relu = nn.ReLU().cuda(device)
def forward(self, x):
x.cuda(device)
output = self.layer1(x) # throws the error
output = self.relu(output)
output = self.layer2(output)
return output
def main():
transform = transforms.Compose([
transforms.ToTensor()
])
mnist_trainset = datasets.MNIST(root='D:\\MNIST', train=True, download=False, transform=transform)
mnist_testset = datasets.MNIST(root='D:\\MNIST', train=False, download=False, transform=transform)
train_loader = DataLoader(dataset=mnist_trainset, batch_size=100, shuffle=True)
test_loader = DataLoader(dataset=mnist_testset, batch_size=100, shuffle=False)
input_size = 784
hidden_size = 500
output_size = 10
num_epochs = 5
learning_rate = 0.001
model = NeuralNet(input_size, hidden_size, output_size)
model.cuda(device)
lossFunction = nn.CrossEntropyLoss()
lossFunction.cuda(device)
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)
losses_in_epochs = []
total_step = len(train_loader)
for epoch in range(num_epochs):
for i, (images, labels) in enumerate(train_loader):
images = images.to(device)
labels = labels.to(device)
images = images.reshape(-1, 28 * 28)
out = model(images)
loss = lossFunction(out, labels)
optimizer.zero_grad()
loss.backward()
optimizer.step()
if (i + 1) % 100 == 0:
print('Epoch [{}/{}], Step [{}/{}], Loss: {:.4f}'.format(epoch + 1, num_epochs, i + 1, total_step,
loss.item()))
if (i % 600) == 0:
losses_in_epochs.append(loss.item())
with torch.no_grad():
correct = 0
total = 0
for images, labels in test_loader:
images = images.reshape(-1, 28 * 28)
out = model(images)
_, predicted = torch.max(out.data, 1)
total += labels.size(0)
correct += (predicted == labels).sum().item()
print('Accuracy of the network on the 10000 test images: {} %'.format(100 * correct / total))
if __name__ == '__main__':
main()
Программный стек:
Python 3.7.1
torch 1.0.1 (with Cuda 9.0)
Windows 10 64-bit