这是pytorch官方的一个例子
官方教程地址:http://pytorch.org/tutorials/beginner/blitz/cifar10_tutorial.html#sphx-glr-beginner-blitz-cifar10-tutorial-py
代码如下
1 # coding=utf-8 2 import torch.nn as nn 3 import torch.nn.functional as F 4 from torch.autograd import Variable 5 import torch 6 import torchvision 7 import torchvision.transforms as transforms 8 import torch.optim as optim 9 10 # The output of torchvision datasets are PILImage images of range [0, 1]. 11 # We transform them to Tensors of normalized range [-1, 1] 12 transform = transforms.Compose([transforms.ToTensor(), 13 transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)), 14 ]) 15 16 # 训练集,将相对目录./data下的cifar-10-batches-py文件夹中的全部数据(50000张图片作为训练数据)加载到内存中,若download为True时,会自动从网上下载数据并解压 17 trainset = torchvision.datasets.CIFAR10(root=r'E:\Face Recognition\cifar-10-python', train=True, download=False, transform=transform) 18 19 # 将训练集的50000张图片划分成12500份,每份4张图,用于mini-batch输入。shffule=True在表示不同批次的数据遍历时,打乱顺序。num_workers=2表示使用两个子进程来加载数据 20 trainloader = torch.utils.data.DataLoader(trainset, batch_size=4, 21 shuffle=True) 22 23 # 测试集,将相对目录./data下的cifar-10-batches-py文件夹中的全部数据(10000张图片作为测试数据)加载到内存中,若download为True时,会自动从网上下载数据并解压 24 testset = torchvision.datasets.CIFAR10(root=r'E:\Face Recognition\cifar-10-python', train=False, download=False, transform=transform) 25 26 # 将测试集的10000张图片划分成2500份,每份4张图,用于mini-batch输入。 27 testloader = torch.utils.data.DataLoader(testset, batch_size=4, 28 shuffle=False) 29 classes = ('plane', 'car', 'bird', 'cat', 30 'deer', 'dog', 'frog', 'horse', 'ship', 'truck') 31 32 33 class Net(nn.Module): 34 def __init__(self): 35 super(Net, self).__init__() 36 self.conv1 = nn.Conv2d(3, 6, 5) # 定义conv1函数的是图像卷积函数:输入为图像(3个频道,即彩色图),输出为6张特征图, 卷积核为5x5正方形 37 self.pool = nn.MaxPool2d(2, 2) 38 self.conv2 = nn.Conv2d(6, 16, 5) 39 self.fc1 = nn.Linear(16 * 5 * 5, 120) 40 self.fc2 = nn.Linear(120, 84) 41 self.fc3 = nn.Linear(84, 10) 42 43 def forward(self, x): 44 x = self.pool(F.relu(self.conv1(x))) 45 x = self.pool(F.relu(self.conv2(x))) 46 x = x.view(-1, 16 * 5 * 5) 47 x = F.relu(self.fc1(x)) 48 x = F.relu(self.fc2(x)) 49 x = self.fc3(x) 50 return x 51 52 53 net = Net() 54 55 criterion = nn.CrossEntropyLoss() # 叉熵损失函数 56 optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9) # 使用SGD(随机梯度下降)优化,学习率为0.001,动量为0.9 57 58 for epoch in range(10): # 遍历数据集两次 59 60 running_loss = 0.0 61 # enumerate(sequence, [start=0]),i序号,data是数据 62 for i, data in enumerate(trainloader, 0): 63 # get the inputs 64 inputs, labels = data # data的结构是:[4x3x32x32的张量,长度4的张量] 65 66 # wrap them in Variable 67 inputs, labels = Variable(inputs), Variable(labels) # 把input数据从tensor转为variable 68 69 # zero the parameter gradients 70 optimizer.zero_grad() # 将参数的grad值初始化为0 71 72 # forward + backward + optimize 73 outputs = net(inputs) 74 loss = criterion(outputs, labels) # 将output和labels使用叉熵计算损失 75 loss.backward() # 反向传播 76 optimizer.step() # 用SGD更新参数 77 78 # 每2000批数据打印一次平均loss值 79 running_loss += loss.data[0] # loss本身为Variable类型,所以要使用data获取其Tensor,因为其为标量,所以取0 80 if i % 2000 == 1999: # 每2000批打印一次 81 print('[%d, %5d] loss: %.3f' % (epoch + 1, i + 1, running_loss / 2000)) 82 running_loss = 0.0 83 84 print('Finished Training') 85 86 correct = 0 87 total = 0 88 for data in testloader: 89 images, labels = data 90 outputs = net(Variable(images)) 91 # print outputs.data 92 _, predicted = torch.max(outputs.data, 1) # outputs.data是一个4x10张量,将每一行的最大的那一列的值和序号各自组成一个一维张量返回,第一个是值的张量,第二个是序号的张量。 93 total += labels.size(0) 94 correct += (predicted == labels).sum() # 两个一维张量逐行对比,相同的行记为1,不同的行记为0,再利用sum(),求总和,得到相同的个数。 95 96 print('Accuracy of the network on the 10000 test images: %d %%' % (100 * correct / total))