首先贴一份在cpu上运行的代码
1 import torch 2 from torchvision import transforms 3 from torchvision import datasets 4 from torch.utils.data import DataLoader 5 import torch.nn.functional as F 6 import torch.optim as optim 7 8 batch_size = 64 9 transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))]) 10 train_dataset = datasets.MNIST(root='../dataset/mnist/', train=True, download=True, transform=transform) 11 train_loader = DataLoader(train_dataset, shuffle=True, batch_size=batch_size) 12 test_dataset = datasets.MNIST(root='../dataset/mnist/', train=False, download=True, transform=transform) 13 test_loader = DataLoader(test_dataset, shuffle=False, batch_size=batch_size) 14 15 16 class Net(torch.nn.Module): 17 def __init__(self): 18 super(Net, self).__init__() 19 self.l1 = torch.nn.Linear(784, 512) 20 self.l2 = torch.nn.Linear(512, 256) 21 self.l3 = torch.nn.Linear(256, 128) 22 self.l4 = torch.nn.Linear(128, 64) 23 self.l5 = torch.nn.Linear(64, 10) 24 25 def forward(self, x): 26 x = x.view(-1, 784) 27 x = F.relu(self.l1(x)) 28 x = F.relu(self.l2(x)) 29 x = F.relu(self.l3(x)) 30 x = F.relu(self.l4(x)) 31 return self.l5(x) 32 33 34 model = Net() 35 36 criterion = torch.nn.CrossEntropyLoss() 37 optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.5) 38 39 40 def train(epoch): 41 running_loss = 0.0 42 for batch_idx, data in enumerate(train_loader, 0): 43 inputs, target = data 44 optimizer.zero_grad() 45 # forward + backward + update 46 outputs = model(inputs) 47 loss = criterion(outputs, target) 48 loss.backward() 49 optimizer.step() 50 running_loss += loss.item() 51 if batch_idx % 300 == 299: 52 print('[%d, %5d] loss: %.3f' % (epoch + 1, batch_idx + 1, running_loss / 300)) 53 running_loss = 0.0 54 55 56 def test(): 57 correct = 0 58 total = 0 59 with torch.no_grad(): 60 for data in test_loader: 61 images, labels = data 62 outputs = model(images) 63 _, predicted = torch.max(outputs.data, dim=1) 64 total += labels.size(0) 65 correct += (predicted == labels).sum().item() 66 print('Accuracy on test set: %d %%' % (100 * correct / total)) 67 68 69 if __name__ == '__main__': 70 for epoch in range(10): 71 train(epoch) 72 test()
要在GPU上运行数据需要把一些相关的参数和模型转到GPU上
需要转换的有：model，数据，criterion（loss函数）
其中optimizer不需要转换
首先定义
1 device = t.device('cuda:0')
将model和criterion to(device)
1 #cuda 2 model = model.to(device) 3 criterion = criterion.to(device)
再将43行的inputs、target，46行的outputs to(device)到GPU上训练
1 def train(epoch): 2 running_loss = 0.0 3 for batch_idx, data in enumerate(train_loader, 0): 4 inputs, target = data 5 #cuda inputs and target 6 inputs = inputs.to(device) ### 7 target = target.to(device) ### 8 optimizer.zero_grad() # noticed: the grad return to zero before starting the loop 9 10 # forward + backward +update 11 outputs = model(inputs) 12 outputs = outputs.to(device) ### 13 loss = criterion(outputs, target) 14 loss.backward() 15 optimizer.step() 16 17 running_loss += loss.item() 18 if batch_idx % 300 == 299: 19 print('[%d, %5d] loss: %.3f' % (epoch + 1, batch_idx + 1, running_loss / 300)) 20 running_loss = 0.0
再将61行的images、labels，62行的outputs to(device)到GPU上测试
1 def test(): 2 correct = 0 3 total = 0 4 with t.no_grad(): # ensuring grad can not updating 5 for data in test_loader: 6 images, label = data 7 #cuda images, label 8 images = images.to(device) ### 9 label = label.to(device) ### 10 outputs = model(images) 11 outputs = outputs.to(device) ### 12 _, predicted = t.max(outputs.data, dim=1) # taking the max num from one_hot code 13 total += label.size(0) 14 # label's format in array that N*1 like [1,8,5,8,3,2,1,4,.....,5]. so label.size() return a array [N,1] and label.size(0) is N 15 correct += (predicted == label).sum().item() 16 print('Accurary on test set: %d %%' % (100 * correct / total))
全部代码（需要cuda而修改或者添加的地方用###在代码后标出）
1 import torch as t 2 from torchvision import transforms 3 from torchvision import datasets 4 from torch.utils.data import DataLoader 5 import torch.optim as op 6 import torch.nn.functional as f 7 8 use_gpu = t.cuda.is_available() ### 9 device = t.device('cuda:0') ### 10 11 batch_size = 64 12 transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))]) 13 # notice : Compose([]) there is []!!!!!!!!!!!!!!!!!!!!!!!! 14 # mnist mean : 0.1307, std :0.3081 15 # convert the PIL to Tensor and then normalized so that the gray level is [0,1] 16 17 train_dataset = datasets.MNIST(root='../dataset/mnist', train=True, download=True, transform=transform) 18 train_loader = DataLoader(train_dataset, batch_size=batch_size) 19 test_dataset = datasets.MNIST(root='../dataset/mnist', train=False, download=True, transform=transform) 20 test_loader = DataLoader(test_dataset, batch_size=batch_size) 21 22 test_dataset = test_dataset 23 train_loader = train_loader 24 test_dataset = test_dataset 25 test_loader = test_loader 26 27 28 class Net(t.nn.Module): 29 def __init__(self): 30 super(Net, self).__init__() 31 self.linear1 = t.nn.Linear(784, 512) 32 self.linear2 = t.nn.Linear(512, 256) 33 self.linear3 = t.nn.Linear(256, 128) 34 self.linear4 = t.nn.Linear(128, 64) 35 self.linear5 = t.nn.Linear(64, 10) 36 37 def forward(self, x): 38 x = x.view(-1, 784) # 1*28*28 Tensor to 1*784 vectors, -1 can auto_transform to pic nums 39 x = f.relu(self.linear1(x)) 40 x = f.relu(self.linear2(x)) 41 x = f.relu(self.linear3(x)) 42 x = f.relu(self.linear4(x)) 43 return self.linear5(x) 44 45 46 model = Net() # there is ()! 47 criterion = t.nn.CrossEntropyLoss() 48 optimizer = op.SGD(model.parameters(), lr=0.01, momentum=0.5) 49 50 # cuda 51 model = model.to(device) ### 52 criterion = criterion.to(device) ### 53 54 55 def train(epoch): 56 running_loss = 0.0 57 for batch_idx, data in enumerate(train_loader, 0): 58 inputs, target = data 59 # cuda inputs and target 60 inputs = inputs.to(device) ### 61 target = target.to(device) ### 62 optimizer.zero_grad() # noticed: the grad return to zero before starting the loop 63 64 # forward + backward +update 65 outputs = model(inputs) 66 outputs = outputs.to(device) ### 67 loss = criterion(outputs, target) 68 loss.backward() 69 optimizer.step() 70 71 running_loss += loss.item() 72 if batch_idx % 300 == 299: 73 print('[%d, %5d] loss: %.3f' % (epoch + 1, batch_idx + 1, running_loss / 300)) 74 running_loss = 0.0 75 76 77 def test(): 78 correct = 0 79 total = 0 80 with t.no_grad(): # ensuring grad can not updating 81 for data in test_loader: 82 images, label = data 83 # cuda images, label 84 images = images.to(device) ### 85 label = label.to(device) ### 86 outputs = model(images) 87 outputs = outputs.to(device) ### 88 _, predicted = t.max(outputs.data, dim=1) # taking the max num from one_hot code 89 total += label.size(0) 90 # label's format in array that N*1 like [1,8,5,8,3,2,1,4,.....,5]. so label.size() return a array [N,1] and label.size(0) is N 91 correct += (predicted == label).sum().item() 92 print('Accurary on test set: %d %%' % (100 * correct / total)) 93 94 95 if __name__ == '__main__': 96 for epoch in range(10): 97 train(epoch) 98 test()