【图像分类】实战——使用VGG16实现对植物幼苗的分类（pytroch）

绝地灬酷狼 2022-10-05 14:54 240阅读 0赞

**目录**

摘要

新建项目

导入所需要的库

设置全局参数

图像预处理

读取数据

设置模型

设置训练和验证

完整代码

--------------------

# 摘要 #

我们这次运用经典的图像分类模型VGG16，实现对植物幼苗的分类，数据集链接：https://pan.baidu.com/s/1JIczDc7VP-PMBnF71302dA 提取码：rqne ，共有12个类别。下面展示图片的样例。

![watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L2hoaGhoaGhoaGh3d3d3d3d3d3d3_size_16_color_FFFFFF_t_70][]![20210609222519301.png][]

大部分的图像是位深度为24位的图像，有个别的是32位的，所以在处理图像时要做强制转换。在这里有一点要提醒大家，拿到数据集，不要上来就搞算法，先去浏览一下数据集，了解数据集是什么样子的，图片有多少，识别难易程度做个初步的认识。

模型采用VGG，模型的详细介绍参照：[【图像分类】一文学会VGGNet（pytorch）\_AI浩-CSDN博客][VGGNet_pytorch_AI_-CSDN]。

接下来讲讲如何使用VGG实现植物幼苗的分类。

# 新建项目 #

新建一个图像分类的项目，data里面放数据集，dataset文件夹中自定义数据的读取方法，这次我不采用默认的读取方式，太简单没啥意思。然后再新建train.py和test.py

在项目的根目录新建train.py，然后在里面写训练代码。

# 导入所需要的库 #

# 设置全局参数 #

设置BatchSize、学习率和epochs，判断是否有cuda环境，如果没有设置为cpu。

# 设置全局参数
    modellr = 1e-4
    BATCH_SIZE = 32
    EPOCHS = 10
    DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# 图像预处理 #

在做图像与处理时，train数据集的transform和验证集的transform分开做，train的图像处理出了resize和归一化之外，还可以设置图像的增强，比如旋转、随机擦除等一系列的操作，验证集则不需要做图像增强，另外不要盲目的做增强，不合理的增强手段很可能会带来负作用，甚至出现Loss不收敛的情况

# 数据预处理
    
    transform = transforms.Compose([
        transforms.Resize((224, 224)),
        transforms.ToTensor(),
        transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])
    
    ])
    transform_test = transforms.Compose([
        transforms.Resize((224, 224)),
        transforms.ToTensor(),
        transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])
    ])

# 读取数据 #

将数据集解压后放到data文件夹下面，如图：

![watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L2hoaGhoaGhoaGh3d3d3d3d3d3d3_size_16_color_FFFFFF_t_70 1][]

然后我们在dataset文件夹下面新建 \_\_init\_\_.py和dataset.py，在dataset.py文件夹写入下面的代码：

说一下代码的核心逻辑。

第一步 建立字典，定义类别对应的ID,用数字代替类别。

第二步 在\_\_init\_\_里面编写获取图片路径的方法。测试集只有一层路径直接读取，训练集在train文件夹下面是类别文件夹，先获取到类别，再获取到具体的图片路径。然后使用sklearn中切分数据集的方法，按照7：3的比例切分训练集和验证集。

第三步 在\_\_getitem\_\_方法中定义读取单个图片和类别的方法，由于图像中有位深度32位的，所以我在读取图像的时候做了转换。

# coding:utf8
    import os
    from PIL import Image
    from torch.utils import data
    from torchvision import transforms as T
    from sklearn.model_selection import train_test_split
    
    Labels = {'Black-grass': 0, 'Charlock': 1, 'Cleavers': 2, 'Common Chickweed': 3,
              'Common wheat': 4, 'Fat Hen': 5, 'Loose Silky-bent': 6, 'Maize': 7, 'Scentless Mayweed': 8,
              'Shepherds Purse': 9, 'Small-flowered Cranesbill': 10, 'Sugar beet': 11}
    
    
    class SeedlingData (data.Dataset):
    
        def __init__(self, root, transforms=None, train=True, test=False):
            """
            主要目标： 获取所有图片的地址，并根据训练，验证，测试划分数据
            """
            self.test = test
            self.transforms = transforms
    
            if self.test:
                imgs = [os.path.join(root, img) for img in os.listdir(root)]
                self.imgs = imgs
            else:
                imgs_labels = [os.path.join(root, img) for img in os.listdir(root)]
                imgs = []
                for imglable in imgs_labels:
                    for imgname in os.listdir(imglable):
                        imgpath = os.path.join(imglable, imgname)
                        imgs.append(imgpath)
                trainval_files, val_files = train_test_split(imgs, test_size=0.3, random_state=42)
                if train:
                    self.imgs = trainval_files
                else:
                    self.imgs = val_files
    
        def __getitem__(self, index):
            """
            一次返回一张图片的数据
            """
            img_path = self.imgs[index]
            img_path=img_path.replace("\\",'/')
            if self.test:
                label = -1
            else:
                labelname = img_path.split('/')[-2]
                label = Labels[labelname]
            data = Image.open(img_path).convert('RGB')
            data = self.transforms(data)
            return data, label
    
        def __len__(self):
            return len(self.imgs)

然后我们在train.py调用SeedlingData读取数据 ，记着导入刚才写的dataset.py(from dataset.dataset import SeedlingData)

dataset_train = SeedlingData('data/train', transforms=transform, train=True)
    dataset_test = SeedlingData("data/train", transforms=transform_test, train=False)
    # 读取数据
    print(dataset_train.imgs)
    
    # 导入数据
    train_loader = torch.utils.data.DataLoader(dataset_train, batch_size=BATCH_SIZE, shuffle=True)
    test_loader = torch.utils.data.DataLoader(dataset_test, batch_size=BATCH_SIZE, shuffle=False)

# 设置模型 #

使用CrossEntropyLoss作为loss，模型采用alexnet，选用预训练模型。更改全连接层，将最后一层类别设置为12，然后将模型放到DEVICE。优化器选用Adam。

# 实例化模型并且移动到GPU
    criterion = nn.CrossEntropyLoss()
    model_ft = vgg16(pretrained=True)
    model_ft.classifier = classifier = nn.Sequential(
                nn.Linear(512 * 7 * 7, 4096),
                nn.ReLU(True),
                nn.Dropout(),
                nn.Linear(4096, 4096),
                nn.ReLU(True),
                nn.Dropout(),
                nn.Linear(4096, 12),
            )
    model_ft.to(DEVICE)
    # 选择简单暴力的Adam优化器，学习率调低
    optimizer = optim.Adam(model_ft.parameters(), lr=modellr)
    
    
    def adjust_learning_rate(optimizer, epoch):
        """Sets the learning rate to the initial LR decayed by 10 every 30 epochs"""
        modellrnew = modellr * (0.1 ** (epoch // 50))
        print("lr:", modellrnew)
        for param_group in optimizer.param_groups:
            param_group['lr'] = modellrnew

# 设置训练和验证 #

# 定义训练过程
    
    def train(model, device, train_loader, optimizer, epoch):
        model.train()
        sum_loss = 0
        total_num = len(train_loader.dataset)
        print(total_num, len(train_loader))
        for batch_idx, (data, target) in enumerate(train_loader):
            data, target = Variable(data).to(device), Variable(target).to(device)
            output = model(data)
            loss = criterion(output, target)
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            print_loss = loss.data.item()
            sum_loss += print_loss
            if (batch_idx + 1) % 10 == 0:
                print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
                    epoch, (batch_idx + 1) * len(data), len(train_loader.dataset),
                           100. * (batch_idx + 1) / len(train_loader), loss.item()))
        ave_loss = sum_loss / len(train_loader)
        print('epoch:{},loss:{}'.format(epoch, ave_loss))
    
    
    # 验证过程
    def val(model, device, test_loader):
        model.eval()
        test_loss = 0
        correct = 0
        total_num = len(test_loader.dataset)
        print(total_num, len(test_loader))
        with torch.no_grad():
            for data, target in test_loader:
                data, target = Variable(data).to(device), Variable(target).to(device)
                output = model(data)
                loss = criterion(output, target)
                _, pred = torch.max(output.data, 1)
                correct += torch.sum(pred == target)
                print_loss = loss.data.item()
                test_loss += print_loss
            correct = correct.data.item()
            acc = correct / total_num
            avgloss = test_loss / len(test_loader)
            print('\nVal set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n'.format(
                avgloss, correct, len(test_loader.dataset), 100 * acc))
    
    
    # 训练
    
    for epoch in range(1, EPOCHS + 1):
        adjust_learning_rate(optimizer, epoch)
        train(model_ft, DEVICE, train_loader, optimizer, epoch)
        val(model_ft, DEVICE, test_loader)
    torch.save(model_ft, 'model.pth')

测试

我介绍两种常用的测试方式，第一种是通用的，通过自己手动加载数据集然后做预测，具体操作如下：

测试集存放的目录如下图：

第一步 定义类别，这个类别的顺序和训练时的类别顺序对应，一定不要改变顺序！！！！

第二步 定义transforms，transforms和验证集的transforms一样即可，别做数据增强。

第三步 加载model，并将模型放在DEVICE里，

第四步 读取图片并预测图片的类别，在这里注意，读取图片用PIL库的Image。不要用cv2，transforms不支持。

import torch.utils.data.distributed
    import torchvision.transforms as transforms
    from PIL import Image
    from torch.autograd import Variable
    import os
    classes = ('Black-grass', 'Charlock', 'Cleavers', 'Common Chickweed',
               'Common wheat','Fat Hen', 'Loose Silky-bent',
               'Maize','Scentless Mayweed','Shepherds Purse','Small-flowered Cranesbill','Sugar beet')
    transform_test = transforms.Compose([
             transforms.Resize((224, 224)),
            transforms.ToTensor(),
            transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])
    ])
    
    DEVICE = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    model = torch.load("model.pth")
    model.eval()
    model.to(DEVICE)
    
    path='data/test/'
    testList=os.listdir(path)
    for file in testList:
            img=Image.open(path+file)
            img=transform_test(img)
            img.unsqueeze_(0)
            img = Variable(img).to(DEVICE)
            out=model(img)
            # Predict
            _, pred = torch.max(out.data, 1)
            print('Image Name:{},predict:{}'.format(file,classes[pred.data.item()]))

第二种 使用自定义的Dataset读取图片

import torch.utils.data.distributed
    import torchvision.transforms as transforms
    from dataset.dataset import SeedlingData
    from torch.autograd import Variable
    
    classes = ('Black-grass', 'Charlock', 'Cleavers', 'Common Chickweed',
               'Common wheat','Fat Hen', 'Loose Silky-bent',
               'Maize','Scentless Mayweed','Shepherds Purse','Small-flowered Cranesbill','Sugar beet')
    transform_test = transforms.Compose([
        transforms.Resize((224, 224)),
        transforms.ToTensor(),
        transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])
    ])
    
    DEVICE = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    model = torch.load("model.pth")
    model.eval()
    model.to(DEVICE)
    
    dataset_test =SeedlingData('data/test/', transform_test,test=True)
    print(len(dataset_test))
    # 对应文件夹的label
    
    for index in range(len(dataset_test)):
        item = dataset_test[index]
        img, label = item
        img.unsqueeze_(0)
        data = Variable(img).to(DEVICE)
        output = model(data)
        _, pred = torch.max(output.data, 1)
        print('Image Name:{},predict:{}'.format(dataset_test.imgs[index], classes[pred.data.item()]))
        index += 1

# 完整代码 #

train.py

import torch.optim as optim
    import torch
    import torch.nn as nn
    import torch.nn.parallel
    import torch.utils.data
    import torch.utils.data.distributed
    import torchvision.transforms as transforms
    from dataset.dataset import SeedlingData
    from torch.autograd import Variable
    from torchvision.models import vgg16
    
    # 设置全局参数
    modellr = 1e-4
    BATCH_SIZE = 32
    EPOCHS = 10
    DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    
    # 数据预处理
    
    transform = transforms.Compose([
        transforms.Resize((224, 224)),
        transforms.ToTensor(),
        transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])
    
    ])
    transform_test = transforms.Compose([
        transforms.Resize((224, 224)),
        transforms.ToTensor(),
        transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])
    ])
    dataset_train = SeedlingData('data/train', transforms=transform, train=True)
    dataset_test = SeedlingData("data/train", transforms=transform_test, train=False)
    # 读取数据
    print(dataset_train.imgs)
    
    # 导入数据
    train_loader = torch.utils.data.DataLoader(dataset_train, batch_size=BATCH_SIZE, shuffle=True)
    test_loader = torch.utils.data.DataLoader(dataset_test, batch_size=BATCH_SIZE, shuffle=False)
    
    # 实例化模型并且移动到GPU
    criterion = nn.CrossEntropyLoss()
    model_ft = vgg16(pretrained=True)
    model_ft.classifier = classifier = nn.Sequential(
                nn.Linear(512 * 7 * 7, 4096),
                nn.ReLU(True),
                nn.Dropout(),
                nn.Linear(4096, 4096),
                nn.ReLU(True),
                nn.Dropout(),
                nn.Linear(4096, 12),
            )
    model_ft.to(DEVICE)
    # 选择简单暴力的Adam优化器，学习率调低
    optimizer = optim.Adam(model_ft.parameters(), lr=modellr)
    
    
    def adjust_learning_rate(optimizer, epoch):
        """Sets the learning rate to the initial LR decayed by 10 every 30 epochs"""
        modellrnew = modellr * (0.1 ** (epoch // 50))
        print("lr:", modellrnew)
        for param_group in optimizer.param_groups:
            param_group['lr'] = modellrnew
    
    
    # 定义训练过程
    
    def train(model, device, train_loader, optimizer, epoch):
        model.train()
        sum_loss = 0
        total_num = len(train_loader.dataset)
        print(total_num, len(train_loader))
        for batch_idx, (data, target) in enumerate(train_loader):
            data, target = Variable(data).to(device), Variable(target).to(device)
            output = model(data)
            loss = criterion(output, target)
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            print_loss = loss.data.item()
            sum_loss += print_loss
            if (batch_idx + 1) % 10 == 0:
                print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
                    epoch, (batch_idx + 1) * len(data), len(train_loader.dataset),
                           100. * (batch_idx + 1) / len(train_loader), loss.item()))
        ave_loss = sum_loss / len(train_loader)
        print('epoch:{},loss:{}'.format(epoch, ave_loss))
    
    
    # 验证过程
    def val(model, device, test_loader):
        model.eval()
        test_loss = 0
        correct = 0
        total_num = len(test_loader.dataset)
        print(total_num, len(test_loader))
        with torch.no_grad():
            for data, target in test_loader:
                data, target = Variable(data).to(device), Variable(target).to(device)
                output = model(data)
                loss = criterion(output, target)
                _, pred = torch.max(output.data, 1)
                correct += torch.sum(pred == target)
                print_loss = loss.data.item()
                test_loss += print_loss
            correct = correct.data.item()
            acc = correct / total_num
            avgloss = test_loss / len(test_loader)
            print('\nVal set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n'.format(
                avgloss, correct, len(test_loader.dataset), 100 * acc))
    
    
    # 训练
    
    for epoch in range(1, EPOCHS + 1):
        adjust_learning_rate(optimizer, epoch)
        train(model_ft, DEVICE, train_loader, optimizer, epoch)
        val(model_ft, DEVICE, test_loader)
    torch.save(model_ft, 'model.pth')

[vgg实现植物幼苗分类.rar-深度学习文档类资源-CSDN下载][vgg_.rar-_-CSDN]

[watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L2hoaGhoaGhoaGh3d3d3d3d3d3d3_size_16_color_FFFFFF_t_70]: https://img-blog.csdnimg.cn/20210609222456647.png?x-oss-process=image/watermark,type_ZmFuZ3poZW5naGVpdGk,shadow_10,text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L2hoaGhoaGhoaGh3d3d3d3d3d3d3,size_16,color_FFFFFF,t_70
[20210609222519301.png]: /images/20221005/34b27d8e635049539d2666d2bd4b92f3.png
[VGGNet_pytorch_AI_-CSDN]: https://blog.csdn.net/hhhhhhhhhhwwwwwwwwww/article/details/117698392?spm=1001.2014.3001.5501
[watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L2hoaGhoaGhoaGh3d3d3d3d3d3d3_size_16_color_FFFFFF_t_70 1]: /images/20221005/568ccce4c32343d69507d8f0d82e0f6a.png
[vgg_.rar-_-CSDN]: https://download.csdn.net/download/hhhhhhhhhhwwwwwwwwww/19526365