图像识别实现流程图

加载图像数据集

  1. #需要用到的库
  2. import torchvision
  3. import torchvision.transforms as transforms
  4. from torch.utils.data import DataLoader
  6. train_batch_size = 4
  7. test_batch_size = 4
  8. num_workers = 0 #线程数
  10. #加载数据集
  11. transform = transforms.Compose([transforms.ToTensor(),transforms.Normalize([0.5,0.5,0.5],[0.5,0.5,0.5])]) #用于数据增强
  12. train_dataset = torchvision.datasets.CIFAR10('./dataset',train=True,transform=transform,download=True)
  13. test_dataset = torchvision.datasets.CIFAR10('./dataset',train=False,transform=transform,download=True)
  14. train_loader = DataLoader(train_dataset,batch_size=train_batch_size,shuffle=True,num_workers=num_workers)
  15. test_loader = DataLoader(test_dataset,batch_size=test_batch_size,shuffle=False,num_workers=num_workers)

完整代码

  1. import torch
  2. import torchvision
  3. import torchvision.transforms as transforms
  4. from torch.utils.data import DataLoader
  6. train_batch_size = 4
  7. test_batch_size = 4
  8. num_workers = 0 #线程数
  9. classes = ('plane','car','bird','cat','deer','dog','frog','horse','ship','truck')
  10. lr = 0.001
  11. momentum = 0.9
  13. #加载数据集
  14. transform = transforms.Compose([transforms.ToTensor(),transforms.Normalize([0.5,0.5,0.5],[0.5,0.5,0.5])])
  15. train_dataset = torchvision.datasets.CIFAR10('./dataset',train=True,transform=transform,download=True)
  16. test_dataset = torchvision.datasets.CIFAR10('./dataset',train=False,transform=transform,download=True)
  17. train_loader = DataLoader(train_dataset,batch_size=train_batch_size,shuffle=True,num_workers=num_workers)
  18. test_loader = DataLoader(test_dataset,batch_size=test_batch_size,shuffle=False,num_workers=num_workers)
  20. #数据可视化
  21. import matplotlib.pyplot as plt
  22. import numpy as np
  23. plt.figure()
  24. def imshow(img):
  25.     img = img/2 +0.5
  26.     npimg = img.numpy()
  27.     plt.imshow(np.transpose(npimg,(1,2,0)))
  28.     plt.show()
  29. examples = enumerate(train_loader)
  30. idx, (examples_data, examples_target) = next(examples) #examples_target是标签列表,0-9表示不同的类别
  31. imshow(torchvision.utils.make_grid(examples_data))
  32. #用于具体查看examples
  33. print('--------------测试examples------------')
  34. print('examples_target.shape:{}'.format(examples_target.shape))
  35. print('examples_target[0]:{}'.format(examples_target[0]))
  36. print('examples_data.shape:{}'.format(examples_data.shape))
  39. #构建网络
  40. import torch.nn as nn
  41. import torch.nn.functional as F
  42. device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
  43. class CNNNet(nn.Module):
  44.     def __init__(self):
  45.         super(CNNNet, self).__init__()
  46.         self.conv1 = nn.Conv2d(in_channels=3,out_channels=16,kernel_size=5,stride=1)
  47.         self.pool1 = nn.MaxPool2d(kernel_size=2,stride=2)
  48.         self.conv2 = nn.Conv2d(in_channels=16,out_channels=36,kernel_size=3,stride=1)
  49.         self.pool2 = nn.MaxPool2d(kernel_size=2,stride=2)
  50.         #self.aap = nn.AdaptiveAvgPool2d(1)
  51.         self.fc1 = nn.Linear(1296,128)
  52.         self.fc2 = nn.Linear(128,10)
  53.         #self.fc3 = nn.Linear(36,10)
  54.     def forward(self,x):
  55.         x = self.pool1(F.relu(self.conv1(x)))
  56.         x = self.pool2(F.relu(self.conv2(x)))
  57.         #x = self.aap(x)
  58.         #x = x.view(x.shape[0],-1)
  59.         #x = self.fc3(x)
  60.         x = x.view(-1,36*6*6)
  61.         #print("x.shape:{}".format(x.shape))
  62.         x = F.relu(self.fc2(F.relu(self.fc1(x))))
  63.         return x
  65. model = CNNNet()
  66. model = model.to(device)
  67. print('--------------查看网络结构-----------')
  68. print(model)
  70. #--训练模型--
  71. print('-----训练优化器-------')
  72. import torch.optim as optim
  73. criterion = nn.CrossEntropyLoss()
  74. optimizer = optim.SGD(model.parameters(), lr=lr, momentum=momentum)
  75. print("----------正式训练模型---------")
  76. losses = []
  77. acces = []
  78. eval_losses = []
  79. eval_acces = []
  80. for epoch in range(10):
  81.     train_acc = 0
  82.     train_loss = 0
  83.     num_correct = 0
  84.     model.train()
  85.     for i, data in enumerate(train_loader):
  86.         img, label = data
  87.         img, label = img.to(device), label.to(device)
  88.         #权重参数梯度清零
  89.         optimizer.zero_grad()
  90.         #正向反向传播
  91.         out = model(img)
  92.         loss = criterion(out, label)
  93.         loss.backward()
  94.         optimizer.step()
  96.         #计算损失值
  97.         train_loss += loss.item()
  98.         #计算准确率
  99.         _, pred = out.max(1)
  100.         num_correct += (pred == label).sum()
  101.         if i % 2000 == 1999:
  102.             print('[%d,%5d] loss : %.3f' % (epoch + 1, i + 1, train_loss / 2000))
  103.             train_loss = 0.0
  104.     acces.append(num_correct/(len(train_loader)*train_batch_size))
  105. #精确率可视化
  106. plt.title('Train Acc')
  107. plt.plot(np.arange(len(acces)),acces)
  108. plt.legend(['Train Acc'],loc='upper right')
  109. plt.show()
  110. #测试模型
  111. eval_loss = 0
  112. eval_acc = 0
  113. class_correct = list(0. for i in range(10))
  114. class_total = list(0. for i in range(10))
  115. total = 0
  116. model.eval()
  117. with torch.no_grad():
  118.     for img, label in test_loader:
  119.         img, label = img.to(device), label.to(device)
  120.         out = model(img)
  121.         #计算损失值
  122.         loss = criterion(out,label)
  123.         eval_loss += loss.item()
  124.         #计算准确率
  125.         _, pred = out.max(1)
  126.         #print("len(label):{}".format(len(label)))
  127.         num_correct += (pred == label).sum()
  128.         c = (pred == label).squeeze()
  129.         acc = num_correct/len(label)
  130.         eval_acc += acc
  131.         total += label.size(0)
  132.         #计算各类别准确率
  133.         for i in range(4):
  134.             class_correct[label[i]] += c[i].item()
  135.             class_total[label[i]] += 1
  136.     eval_losses.append(eval_loss/total)
  137.     eval_acces.append(eval_acc/total)
  138.     print("total:{}".format(total))
  139.     print("len(test_loader):{}".format(len(test_loader)))
  140.     for i in range(10):
  141.         print("accuracy of {}:{}%".format(classes[i],100*class_correct[i]/class_total[i]))
  142.     print("----------------")
  143.     #print('epoch:{}, eval_loss:{:.4f},eval_acc:{:.4f}'.format(epoch,eval_loss/len(test_loader),eval_acc/len(test_loader)))
  144.     #rint("Accuracy of the network on the 10000 test images:%d %%" % (100 * eval_acc / len(test_loader)))

输出结果

image.png

部分输出结果解释

[epoch, i] （如[10,8000])，10代表的是进行到的迭代次数，8000指的是dataloader进行到了第8000份（每2000份打印一回），由于train_loader一共有12500份，因此最多只能打印到12000份，不会打印到14000份。但是实际上会进行到12500份的。