前言

        图像分类这里对小企业还是特定的工业企业都有着非常重要的意义,不是简简单单的那一个大模型来糊弄的,是需要根据特定的环境来进行识别的,我们经常能在网上看到农业上用于杂草以及病虫的识别,之后使用激光进行精准打击,这就使用的是图像识别,但是往往图像识别需要使用到的环境内容都比较庞大今天我来使用AlexNet来完成一个图像分类识别的模型炼制,前期的图像处理可以使用本地的电脑来处理,毕竟只是图片文件的处理,后面的pytorch来完成的模型训练需要使用到GPU来处理,我本地的只有8G肯定是不行的,我在丹摩上选了一个24G显存的来炼制,看看效率如何。

目录

前言

正文

1、图像数据制作与处理

图像尺寸预处理

图像统一命名

图像统一后缀名

测试图像

划分训练集和测试集

2、训练

丹摩服务器购买

上传数据集

mean与std值计算

运行train.py进行训练

3、测试

总结


正文

这里我们分两个大步骤来进行,再加上测试步骤,共计3步,先进行图像的处理,也就是数据集的处理,第二步是根据数据集进行云联,最后我们再测试一下。

1、图像数据制作与处理

图像我这里就找了一些网上的网红图片,截取了几十张,主要是两个人,一个是穗岁,一个是小雨

图像尺寸预处理

图像大小统一处理的话我建议使用stablediffusion的训练图像预处理啊,一键搞定,很方便。

注:

AlexNet是一个深度学习模型,它是由Alex Krizhevsky、Ilya Sutskever和Geoffrey Hinton在2012年提出的。AlexNet采用了8层(包括5个卷积层和2个全连接层)的结构,其中卷积层后面跟着最大池化层,并在每个卷积层和全连接层之后使用了修剪(ReLU)激活函数,以及在全连接层之后使用了Dropout来防止过拟合。

在AlexNet模型中,输入图像的尺寸是固定的,通常是224x224(对于原始的AlexNet模型)或227x227(对于后来的一些改进版本),颜色通道数目是3(RGB)。

一键预处理完毕。

图像统一命名

图像大小处理完毕需要处理图像的统一命名,这样是为了更好的分类。

import os

list_dir = os.listdir('dataset')


for dir_name in list_dir:
    image_list = os.listdir("dataset/" + dir_name)
    for i in range(len(image_list)):
        print(i)
        img_path = ("dataset\\" + dir_name + "\\" + image_list[i])
        image_name_new = dir_name + "_" + str(i) + "." + image_list[i].split('.')[-1]
        image_name_new_path = "dataset/" + dir_name + "/" + image_name_new
        print(image_name_new_path)
        os.rename(img_path, image_name_new_path)

图像统一后缀名

名称统一完毕后,我们再讲图片的后缀统一一下。

import os
import glob
from PIL import Image

image_list = glob.glob('dataset/*/*.*')
for image_name in image_list:
    imag_suffix = image_name.split('.')[-1]
    if imag_suffix != 'jpg':
        im = Image.open(image_name)
        im = im.convert('RGB')
        image_name_new = image_name.split('.')[-2] + ".jpg"
        print(image_name_new)
        im.save(image_name_new, quality=95)
        os.remove(image_name)

测试图像

随机选出几个用作最后的测试。

import shutil
import glob
image_list = glob.glob('dataset/*/*.*')
from sklearn.model_selection import train_test_split
import os
os.makedirs("test",exist_ok=True)
trainval_files, test_files = train_test_split(image_list, test_size=0.05, random_state=42)

for file in test_files:
    file=file.replace("\\","/")
    file_name=file.split("/")[-1]
    print(file_name)
    file_new_path="test/"+file_name
    shutil.move(file,file_new_path)

划分训练集和测试集

这里就是先将所有的图片都存储到list当中,再通过sklearn.model_selection的train_test_split方法将list拆分。缺少库的话自己ALT+回车安装即可。

import glob
import os
import shutil
from sklearn.model_selection import train_test_split

image_list=glob.glob('dataset/*/*.*')
print(image_list)
file_dir='data'
if os.path.exists(file_dir):
    print('true')
    shutil.rmtree(file_dir)#删除再建立
    os.makedirs(file_dir)
else:
    os.makedirs(file_dir)
trainval_files, val_files = train_test_split(image_list, test_size=0.2, random_state=42)
train_dir='train'
val_dir='val'
train_root=os.path.join(file_dir,train_dir)
val_root=os.path.join(file_dir,val_dir)
for file in trainval_files:
    file_class=file.replace("\\","/").split('/')[-2]
    file_name=file.replace("\\","/").split('/')[-1]
    file_class=os.path.join(train_root,file_class)
    if not os.path.isdir(file_class):
        os.makedirs(file_class)
    shutil.copy(file, file_class + '/' + file_name)

for file in val_files:
    file_class=file.replace("\\","/").split('/')[-2]
    file_name=file.replace("\\","/").split('/')[-1]
    file_class=os.path.join(val_root,file_class)
    if not os.path.isdir(file_class):
        os.makedirs(file_class)
    shutil.copy(file, file_class + '/' + file_name)

2、训练

训练这里操作的内容相对来说是比较方便的,而且在我们24G显存的作用下可以看到速度的非常不错的呢,所以我们得先购买一下丹摩的服务器。

丹摩服务器购买

丹摩的官网首页:

丹摩DAMODEL|让AI开发更简单!算力租赁上丹摩!

注册完毕后认证一下,然后直接进来创建实例就行。

因为认证完有50的代金券呢。

磁盘大小一定大一些啊,默认50肯定不够用。 

不够的话就根据情况再冲就行了。

我们选择PyTorch的2.3.0版本,创建的时候需要个密钥对,创建一下啊,很直接的。

创建中,别着急。

创建完毕我们直接访问一下。

非常空的什么都没有。

上传数据集

上传就是创建文件夹与点击上传按钮了,可以批量文件上传,但是不能有文件夹哦。

我的服务器没有一些环境,我们自己安装一下,例如运行makedata的时候就需要scikit的包,下图有安装提示。

mean与std值计算

from torchvision.datasets import ImageFolder
import torch
from torchvision import transforms

def get_mean_and_std(train_data):
    train_loader = torch.utils.data.DataLoader(
        train_data, batch_size=1, shuffle=False, num_workers=0,
        pin_memory=True)
    mean = torch.zeros(3)
    std = torch.zeros(3)
    for X, _ in train_loader:
        for d in range(3):
            mean[d] += X[:, d, :, :].mean()
            std[d] += X[:, d, :, :].std()
    mean.div_(len(train_data))
    std.div_(len(train_data))
    return list(mean.numpy()), list(std.numpy())

if __name__ == '__main__':
    train_dataset = ImageFolder(root=r'data', transform=transforms.ToTensor())
    print(get_mean_and_std(train_dataset))

可以看到已经计算得到了mean与std值,生成的是两个list,第一个是mean的结果,第二个是std的结果。

这个值好多地方都要用到哦。

mean=[0.5740846, 0.5032924, 0.47821605], std=[0.3002571, 0.28379133, 0.27508244]

运行train.py进行训练

import torch.optim as optim
import torch
import torch.nn as nn
import torch.nn.parallel
import torch.optim
import torch.utils.data
import torch.utils.data.distributed
import torchvision.transforms as transforms
from torchvision.models import alexnet
import os
from torchvision import datasets
import json


def seed_everything(seed=42):
    '''
    随机因子
    '''
    os.environ['PYHTONHASHSEED'] = str(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.backends.cudnn.deterministic = True

# 训练
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 = data.to(device,non_blocking=True), target.to(device,non_blocking=True)
        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,Best_ACC):
    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 = data.to(device,non_blocking=True),target.to(device,non_blocking=True)
            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))
    return acc


if __name__ == '__main__':
    #创建保存模型的文件夹
    file_dir = 'checkpoints/AlexNet'
    if os.path.exists(file_dir):
        print('true')
        os.makedirs(file_dir,exist_ok=True)
    else:
        os.makedirs(file_dir)

    # 设置全局参数
    model_lr = 1e-4
    BATCH_SIZE = 24
    EPOCHS = 300
    DEVICE = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
    classes = 2
    # 初始的时候使用None·如果继续训练就改一下最后的那个resume = 'checkpoints/AlexNet/model_300_1.0.pth'
    resume = None
    Best_ACC = 0 #记录最高得分
    start_epoch=1
    SEED=42
    seed_everything(42)

    # 数据预处理7
    transform = transforms.Compose([
        transforms.RandomRotation(10),
        transforms.GaussianBlur(kernel_size=(5, 5), sigma=(0.1, 3.0)),
        transforms.ColorJitter(brightness=0.5, contrast=0.5, saturation=0.5),
        transforms.Resize((224, 224)),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.5560672, 0.48953587, 0.46780446], std=[0.29265925, 0.27421287, 0.26669672])
    
    ])
    transform_test = transforms.Compose([
        transforms.Resize((224, 224)),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.5560672, 0.48953587, 0.46780446], std=[0.29265925, 0.27421287, 0.26669672])
    ])

    # 读取数据
    dataset_train = datasets.ImageFolder('data/train', transform=transform)
    dataset_test = datasets.ImageFolder("data/val", transform=transform_test)
    with open('class.txt', 'w') as file:
        file.write(str(dataset_train.class_to_idx))
    with open('class.json', 'w', encoding='utf-8') as file:
        file.write(json.dumps(dataset_train.class_to_idx))
    # 导入数据
    train_loader = torch.utils.data.DataLoader(dataset_train, batch_size=BATCH_SIZE, pin_memory=True, shuffle=True)
    test_loader = torch.utils.data.DataLoader(dataset_test, batch_size=BATCH_SIZE, pin_memory=True, shuffle=False)
    
  # 实例化模型并且移动到GPU
    criterion = torch.nn.CrossEntropyLoss()
 	# 设置模型
    model_ft = alexnet(pretrained=True)
    
    num_ftrs=model_ft.classifier[6].in_features
    model_ft.classifier[6]=nn.Linear(num_ftrs,classes)
    # num_ftrs = model_ft.head.in_features
    # model_ft.head = nn.Linear(num_ftrs, classes)
    if resume:
        model = torch.load(resume)
        model_ft.load_state_dict(model['state_dict'])
        Best_ACC = model['Best_ACC']
        start_epoch = model['epoch'] + 1
    model_ft.to(DEVICE)
    print(model_ft)

    # 选择简单暴力的Adam优化器,学习率调低
    optimizer = optim.AdamW(model_ft.parameters(),lr=model_lr)
    
    # 训练
    for epoch in range(start_epoch, EPOCHS + 1):
        train(model_ft, DEVICE, train_loader, optimizer, epoch)
        acc=val(model_ft, DEVICE, test_loader,Best_ACC)
        if acc>=Best_ACC:
            Best_ACC=acc
            torch.save(model_ft, file_dir + "/" + 'best.pth')
        state = {
            'epoch': epoch,
            'state_dict': model_ft.state_dict(),
            'Best_ACC': Best_ACC
        }
        filename=file_dir + "/" + 'model_' + str(epoch) + '_' + str(round(acc, 3)) + '.pth'
        torch.save(state, filename)

一直等待到训练完成。 

我们可以看到【checkpoints】中看到pth文件。

3、测试

我这里就拿图片直接测试了一下,效果很直接,分辨能力很强,其实还能继续训练,这就得看具体的数据集的量了。

import torch.utils.data.distributed
import torchvision.transforms as transforms
from PIL import Image
from torch.autograd import Variable
import os

classes = ('小雨','穗岁')
transform_test = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.5560672, 0.48953587, 0.46780446], std=[0.29265925, 0.27421287, 0.26669672])
])

DEVICE = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
model=torch.load('checkpoints/AlexNet/best.pth',map_location='cpu')
model.eval()
model.to(DEVICE)
path = 'test/'
file = '待预测图片9(雨).jpg'
img = Image.open(path + file).convert('RGB')
img = transform_test(img)
img.unsqueeze_(0)
img = Variable(img).to(DEVICE)
out = model(img)
# Predict
_, pred = torch.max(out.data, 1)
print('图片名称:{},预测结果:{}'.format(file, classes[pred.data.item()]))

总结

我们整体的训练时间还是比较短的,而且次数也不多,如果多一些,例如使用32G的显存训练个大几千回合,那就是最终的可以真正识别物品或者某目标的实际模型了,完全可以应用在企业当中,本文使用的是丹摩的服务器,相对来说价格还是比较美丽的,效果整体体验下来也是不错的,需要使用的可以点这里:丹摩DAMODEL|让AI开发更简单!算力租赁上丹摩!

希望本文可以为大家创造一定的价值。 

Logo

尧米是由西云算力与CSDN联合运营的AI算力和模型开源社区品牌,为基于DaModel智算平台的AI应用企业和泛AI开发者提供技术交流与成果转化平台。

更多推荐