【超详细CNN图像分类实践】Pytorch实现宝石分类

目录

一、任务说明

二、任务准备

三、任务内容

四、任务过程 

1.分割数据集 

2.定义CNN

3.加载数据集 

4.训练模型 

5.测试模型 

一、任务说明

数据集来源于百度飞桨AI Studio中的公开数据集。宝石分类数据集共有25种宝石类别，每种类别的图片数目在30~40张左右，共有811张训练图片和4张测试图片（测试图片少的可怜）。

数据集链接：https://aistudio.baidu.com/datasetdetail/162243

本次任务的实践环境是：Windows11 （RTX 3050显卡）+ PyCharm 2023.2.3 + PyTorch 2.0.1+cu117 + Python3.8

本次任务适合有一点儿深度学习、pytorch和python基础的同学，还涉及tensorboard的简单使用。主要目的是学会用CNN进行图像分类有哪些步骤，每一步要怎么做，不涉及调参、优化网络等知识。 

二、任务准备

没有安装tensorboard的同学需要在终端运行“pip install tensorboard”安装一下，在PyCharm里新建项目文件，在项目目录里新建一个jupyter notebook文件，然后将解压好的数据集以下面的文件夹结构添加进项目目录。

三、任务内容

该任务包含的有分割数据集、定义CNN、加载数据集、训练模型、测试模型五个步骤 。

四、任务过程 

1.分割数据集 

解压数据集文件后发现，原始的测试集图片数量太少（而且是同一类的）， 因此我将测试集的4张图片塞进训练对应的类别里去了，这样就需要重新划分训练集测试集。本次任务的划分方式是每一类每有15张图片就把该张图片作为测试图片，最终得到了46张不同类别的测试图片，代码如下：

import os
import random

def split_dataset(path):
    #先判断存储图像路径和标签的文件是否存在，存在的话就删掉，以免重复录入。
    if os.path.isfile('train.txt') and os.path.isfile('test.txt'):
        os.remove('./train.txt')
        os.remove('./test.txt')
    else:
        pass
    
    #存储训练图片的路径和标签
    train_split = []
    #存储测试图片的路径和标签
    test_split = []
    #图像标签
    target = 0
    #存储标签对应类别名称的字典
    dataset_details = {}
    for class_num in os.listdir(path):
        count = 0
        dataset_details[target] = class_num
        class_path = os.path.join(path, class_num)
        for img_name in os.listdir(class_path):
            count += 1
            img_path = os.path.join(class_path, img_name)
            #每遍历十五张图片就把这张图作为测试图片
            if count % 15 == 0:
                test_split.append(img_path + '\t%d' % target + '\n')
            else:
                train_split.append(img_path + '\t%d' % target + '\n')
        target += 1
    
    #打乱数据集
    random.shuffle(train_split)
    random.shuffle(test_split)
    
    #将存储列表的内容录入文件里，方便后边加载数据
    with open('./train.txt', mode='a') as f1:
        for train_data in train_split:
            f1.write(train_data)
        f1.close()
    with open('./test.txt', mode='a'