import torch
import torch.nn as nn
import torch.optim as optim
import torchvision
import torchvision.transforms as transforms
import time
transform = transforms.Compose([transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])
traindata = torchvision.datasets.CIFAR10("./data", train=True, transform=transform, download=True)
trainloader = torch.utils.data.DataLoader(traindata, batch_size=4, shuffle=True)
testdata = torchvision.datasets.CIFAR10("./data", train=False, transform=transform, download=True)
testloader = torch.utils.data.DataLoader(testdata, batch_size=4, shuffle=False)
class Net(nn.Module):
def __init__(self):
super(Net, self).__init__()
self.conv = nn.Sequential(
nn.Conv2d(3,6,5),
nn.MaxPool2d(2,2),
nn.Conv2d(6,16,5),
nn.MaxPool2d(2, 2),
)
self.linear = nn.Sequential(
nn.Linear(16 * 5 * 5, 160),
nn.ReLU(),
nn.Linear(160, 64),
nn.ReLU(),
nn.Linear(64, 10),
)
def forward(self, inputs):
conv_output = self.conv(inputs)
output = self.linear(conv_output.view(-1,16 * 5 * 5))
return output
device = torch.device("cuda:0")
start = time.time()
model = Net()
model.to(device)
loss_fc = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)
for epoch in range(10):
for i,data in enumerate(trainloader,1):
input, target = data
input, target = input.to(device),target.to(device)
pred = model(input)
loss = loss_fc(pred, target)
optimizer.zero_grad()
loss.backward()
optimizer.step()
if i%2000 == 0:
print("[{}, {}] loss:{:.4f}".format(epoch+1, i, loss.item()))
print("time taken:{:.3f}".format(time.time()-start))
The console print:
time taken:361.207
But if I train the model on CPU, the console print:
time taken:205.772
I checked the Task Manager when training on GPU,and the GPU Occupancy rate was only 5%-10%。
My GPU is GTX 1060 6GB
Why did it take more time on GPU than on CPU?