1. Problems
When I was practicing using pytorch today, I prepared to use GPU, and the following errors occurred:
RuntimeError: Input type (torch.FloatTensor) and weight type (torch.cuda.FloatTensor) should be the same
2. Code (adjusted and can run correctly)
import torch.optim
import torchvision.datasets
# Preparing the data set
from torch import nn
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter
from time import time
print(torch.cuda.is_available())
train_data = torchvision.datasets.CIFAR10(root="./dataset",
train=True,
transform=torchvision.transforms.ToTensor(),
download=True)
test_data = torchvision.datasets.CIFAR10(root="./dataset",
train=False,
transform=torchvision.transforms.ToTensor(),
download=True)
print("training_set_data_length: %d" % len(train_data))
print("Test set data length: %d" % len(test_data))
# Load data using DataLoader
train_data_loader = DataLoader(train_data, batch_size=64)
test_data_loader = DataLoader(test_data, batch_size=64)
# Build the neural network (in a separate .py file)
class Net(nn.Module):
def __init__(self) -> None:
super(Net, self).__init__()
self.model = nn.Sequential(
nn.Conv2d(3, 32, 5, 1, 2),
nn.MaxPool2d(2, 2),
nn.Conv2d(32, 32, 5, 1, 2),
nn.MaxPool2d(2, 2),
nn.Conv2d(32, 64, 5, 1, 2),
nn.MaxPool2d(2, 2),
nn.Flatten(),
nn.Linear(1024, 64),
nn.Linear(64, 10)
)
def forward(self, x):
x = self.model(x)
return x
# Create a network model
net = Net()
# Only the model, data, and loss function can run on the GPU
# if torch.cuda.is_available():
# net = net.cuda()
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
net.to(device)
# Loss function
loss_fn = nn.CrossEntropyLoss()
# if torch.cuda.is_available():
# loss_fn = loss_fn.cuda()
loss_fn.to(device)
# Optimizer
learning_rate = 1e-2 # 0.01
optimizer = torch.optim.SGD(net.parameters(), lr=learning_rate)
# Set the parameters of the training network
# Record the number of training sessions
total_train_step = 0
# Record the number of training sessions
total_test_step = 0
# Number of training rounds
epoch = 10
start_time = time()
writer = SummaryWriter("./logs/train")
for i in range(epoch):
print("------Round %d of training ------" % (i + 1))
# Training steps
for data in train_data_loader:
imgs, targets = data
if torch.cuda.is_available():
imgs, targets = imgs.cuda(), targets.cuda()
# imgs.to(device)
# targets.to(device)
outputs = net(imgs)
loss = loss_fn(outputs, targets)
# Optimizer optimization model
optimizer.zero_grad()
loss.backward()
optimizer.step()
total_train_step += 1
if total_train_step % 100 == 0:
end_time = time()
print(end_time - start_time)
print("training_step: {}, loss: {}".format(total_train_step, loss.item())) # .item() can convert the tensor type to a number
writer.add_scalar("train_loss", loss.item(), total_train_step)
# Test Steps
total_test_loss = 0
total_accuracy = 0
with torch.no_grad(): # Zero out the gradient when testing. No need to adjust the gradient for optimization
for data in test_data_loader:
imgs, targets = data
if torch.cuda.is_available():
imgs, targets = imgs.cuda(), targets.cuda()
# imgs.to(device)
# targets.to(device)
outputs = net(imgs)
loss = loss_fn(outputs, targets)
total_test_loss += loss.item()
accuracy = (outputs.argmax(1) == targets).sum()
total_accuracy += accuracy
print("loss on the overall test set: {}".format(total_test_loss))
print("Percent correct on the overall test set: {}".format(total_accuracy/len(test_data)))
writer.add_scalar("test_loss", total_test_loss, total_test_step)
writer.add_scalar("test_accuracy", total_accuracy/len(test_data), total_test_step)
total_test_step += 1
# Save the model
# torch.save(net.state_dict(), "model_{}.pth".format(i))
# print("Round {} training model saved".format(i))
writer.close()
3. Solutions
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
for data in train_data_loader:
imgs, targets = data
if torch.cuda.is_available():
imgs, targets = imgs.cuda(), targets.cuda()
# imgs.to(device)
# targets.to(device)
outputs = net(imgs)
In the above code, IMGs and targets cannot use .to(device), so the input type (torch.Floattensor) will appear after use. If it is not GPU type, it can only be used in another way:
if torch.cuda.is_available():
imgs, targets = imgs.cuda(), targets.cuda()
This can solve the problem that the input and weight types do not match.
4. Reference
https://stackoverflow.com/questions/59013109/runtimeerror-input-type-torch-floattensor-and-weight-type-torch-cuda-floatte