基于图像分类的工业应用--Pytorch实现

Pytorch实战项目--新手向

第一章 

第2章时尚物品分类普通神经网络

第3章应用AlexNet实现物品分类

文章目录

• 系列文章目录
• 前言
• 一、pandas是什么？
• 二、使用步骤
  • 1.引入库
  • 2.读入数据
• 总结

---

前言

提示：这里可以添加本文要记录的大概内容：

例如：随着人工智能的不断发展，机器学习这门技术也越来越重要，很多人都开启了学习机器学习，本文就介绍了机器学习的基础内容。

---

提示：以下是本篇文章正文内容，下面案例可供参考

一、创建项目，将不同功能写成不同模块

 

二、创建自己的数据集

继承Pytorch的datasets类，且必须实现__len__和__getitem__两个方法

1. 导包

import os

import torch
import numpy
from torch.utils.data import Dataset,DataLoader
from torchvision import transforms
import cv2 as cv

2.创建类

代码如下：

class SurfaceDataset(Dataset):
    def __init__(self,root_dir):
        #使用迁移学习
        self.transform = transforms.Compose([transforms.ToTensor(),
                                             transforms.Normalize(mean=[0.48,0.456,0.406],std=[0.229,0.224,0.225]),
                                             transforms.Resize((200,200))])#从网上找的，要保证预处理的数据分布相同
        img_files = os.listdir(root_dir)#读取所有的图片路径
        self.defect_types = []
        self.images = []
        #存储文件名和图片对应的label的索引
        for file_name in img_files:
            #Cr_11.bmp 以下划线分割文件名
            defect_class = file_name.split('_')[0]
            defect_index = defect_labels.index(defect_class)#获取对应索引
            self.images.append(os.path.join(root_dir,file_name)) #train/Cr
            self.defect_types.append(defect_index)

    def __len__(self):
        return len(self.images)

    def __getitem__(self, idx):
        image_path = self.images[idx]
        img = cv.imread(image_path)#读取进来的图片格式：BGR
        img = cv.cvtColor(img,cv.COLOR_BGR2RGB)#变RGB
        sample = {'image':self.transform(img),'defect':self.defect_types[idx]}
        #包装成字典，根据字典取值
        return sample

3.主函数运行

if __name__ =='__main__':
    ds = SurfaceDataset('./enu_surface_defect/train')
    print(len(ds))
    print(ds[0]['image'].shape, ds[0]['defect'])
    dl = DataLoader(ds, batch_size=8, shuffle=True, num_workers=8)
    sample = next(iter(dl))
    print(type(sample))
    print(sample['image'].shape)

 

三、创建模型

使用迁移学习

import torch
import torchvision


class SurfaceDecter(torch.nn.Module):
    def __init__(self,num_classes = 1000):
        super(SurfaceDecter, self).__init__()
        self.cnn_layers = torchvision.models.resnet18(pretrained = True)#pretrained = True加载的模型经过预训练
        in_features = self.cnn_layers.fc.in_features
        #fc是层的名字
        print(self.cnn_layers)
        print(in_features)
        #修改最后一层的输出维度
        self.cnn_layers.fc = torch.nn.Linear(in_features,num_classes)

    def forward(self,x):
        #完成前向传播
        out = self.cnn_layers(x)
        return out


if __name__ == '__main__':
    model = SurfaceDecter(num_classes=6)

 

四·训练模型

1. 导包
2. 获取设备名
3. 创建数据集
4. 将类别和对应索引两者关系，保存成json文件
5. 多线程实现下载文件功能
6. 读取数据dataset---dataLoader
7. 创建模型
8. 定义训练过程
9. 模型保存（最好的表现）

   import json
   import os
   import sys
   
   import torch
   import torch.nn as nn
   from torchvision import transforms,datasets
   from tqdm import tqdm#可视化训练过程
   
   from my_dadaset import SurfaceDataset
   from my_dadaset import defect_labels
   from model import SurfaceDecter
   from model import SurfaceDectectResNet
   
   
   def main():
       device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
       print(f'using{device}')
   
   
       train_dataset = SurfaceDataset('./enu_surface_defect/train')
       train_num = len(train_dataset)
   
       class_dict = dict((i,label) for i,label in enumerate(defect_labels))
       json_dict = json.dumps(class_dict,indent=4)
       #indent:如果参数为正整数就将其进行漂亮的格式化打印，如果是复数和空值、默认的情况下就是执行紧凑型打印，如果是字符串用于缩进每个几倍。
       with open('class_indices.json','w') as json_file:
           json_file.write(json_dict)
   
       batch_size = 32
       nw = min([os.cpu_count(),batch_size if batch_size > 1 else o,8])
       print(f'使用{nw}个线程')
   
       train_dataset = SurfaceDataset('./enu_surface_defect/train')
       train_loader = torch.utils.data.DataLoader(train_dataset,
                                                  batch_size = batch_size,
                                                  shuffle = True,
                                                  nuw_workers = nw)
   
       validate_dataset = SurfaceDataset('./enu_surface_defect/train')
       val_num = len(validate_dataset)#样本数量
       validate_loader = torch.utils.data.DataLoader(validate_dataset,
                                                  batch_size=batch_size,
                                                  shuffle=False,
                                                  nuw_workers=nw)
   
       net = SurfaceDectectResNet(num_classes = 6):
       net.to(device)
       loss_fn = nn.CrossEntropyLoss()
       optimizer = torch.optim.Adam(net.parameters(),ls = 0.001)
   
   
       epoches = 10
       save_path = './model_weight.pth'
       best_acc = 0.0
       train_step = len(train_loader)#每次读取的数据是一批，按批次读取，长度就是分成几批，就是训练次数
       for epoch in range(epoches):
           #训练模式
           net.train()
           running_loss = 0.0
   
           train_bar = tqdm(train_loader,file=sys.stdout)
           for step,data in enumerate(train_bar):
               images,labels = data['image'],data['defect']
               optimizer.zero_grad()#反向求导前前清零
               #对应变量也要copy到设备上
               x,y = images.to(device),labels.to(device)
               outputs = net(x)
               loss = loss_fn(outputs,y)
               loss.backward()
               optimizer.step()
   
               running_loss += loss.item()#取出torch的值标量
               #输出什么内容
               train_bar.desc = f'train epoch[{epoch + 1}/{epochs}] loss:{loss:.3f}'
   
           #预测模式
           net.eval()
           acc = 0.0
           with torch.no_grad():
               #不用反向传播
               val_bar = tqdm(validate_loader, file=sys.stdout)
               for val_data in val_bar:
                   val_images, val_labels = val_data['image'], val_data['defect']
                   val_x,val_y = val_images.to(device),val_labels.to(device)
                   outputs = net(val_x)
                   predict_y = torch.max(outputs, dim=1)[1]#概率最大值
                   acc += torch.eq(predict_y, val_y).sum().item() #预测出的是0，1，2，3
           val_accuracy = acc / val_num
           print(f'[epoch {epoch + 1} train_loss: {running_loss / train_steps:.3f},'
                 f'val_accuracy:{val_accuracy:.3f}')
           if val_accuracy > best_acc:
               best_acc = val_accuracy
               torch.save(net.state_dict(), save_path)
   
   if __name__ == '__main__':
       main()

---

 

 

总结

使用迁移学习，只修改最后的输出层，可以加速训练，直接用人家训练好的参数