训练你的第一个图像分类深度学习模型

最新推荐文章于 2025-04-16 11:22:34 发布

一个油桃子

最新推荐文章于 2025-04-16 11:22:34 发布

阅读量270

点赞数 4

文章标签：分类深度学习机器学习 python

本文链接：https://blog.youkuaiyun.com/m0_50807424/article/details/139881946

版权

环境准备

我的电脑配置：WSL2(当然使用win10或win11皆可)
显卡：4060
需要安装的库：fastbook

推荐使用conda 创建一个新的python环境，python版本选择3.10即可
激活该环境，在此环境中使用pip命令安装fastbook库，会为你自动安装需要模型所需要的库
不想使用自己的电脑可以使用Kaggle提供的notebook
看下去，你可以训练你想要训练的任何分类模型！！！

开始

# !pip install -Uqq fastbook  若使用Kaggle中的环境，取消此处注释即可安装所需库
import fastbook
fastbook.setup_book()

from fastbook import *
from fastai.vision.widgets import *

注册一个账号，获得Bing Image Search key

key = os.environ.get('AZURE_SEARCH_KEY', 'Your bing search key')

我这里想训练一个分辨猫和鸡的模型

你可以训练任何你想训练的模型，选择下载你感兴趣的图片就可以了

animal_types = 'cat', 'chicken'
path = Path('animal')

使用bing提供的接口下载图片来训练模型

if not path.exists():
    path.mkdir()
    for o in animal_types:
        dest = (path/o)
        dest.mkdir(exist_ok=True)
        results = search_images_bing(key, f'{o} animal')  # 如果你不是训练动物分类模型的话，请修改这里。
        download_images(dest, urls=results.attrgot('contentUrl'))

查看下载了多少张图片

fns = get_image_files(path)
fns

检测多少不能打开的图片

failed = verify_images(fns)
failed

删除掉

failed.map(Path.unlink)

使用数据装载函数来处理训练时如何操作图片

animals = DataBlock(
    blocks=(ImageBlock, CategoryBlock), 
    get_items=get_image_files, 
    splitter=RandomSplitter(valid_pct=0.2, seed=42),
    get_y=parent_label,
    item_tfms=Resize(128))

dls = animals.dataloaders(path)

查看一下验证集中的图片

dls.valid.show_batch(max_n=4, nrows=1)

对图片施加一些图像增强技术

animals= animals.new(
    item_tfms=RandomResizedCrop(224, min_scale=0.5),
    batch_tfms=aug_transforms())
dls = animals.dataloaders(path)

在resnet18的基础上微调，得到一个性能强大的图像分类模型

learn = vision_learner(dls, resnet18, metrics=error_rate)
learn.fine_tune(8)

查看一下混淆矩阵

interp = ClassificationInterpretation.from_learner(learn)
interp.plot_confusion_matrix()

查看损失值最高的几张图片

interp.plot_top_losses(5, nrows=1)

将训练好的模型文件导出，可使用gradio部署一个简易的操作界面

learn.export('catorchicken.pkl')

此时，你会发现你的当前目录多了catorchicken.pkl 这个文件，后面进行推理时，就不用每次训练一遍了。
装载训练好的模型文件

learn_inf = load_learner('catorchicken.pkl')

使用模型进行推理任务

learn_inf.predict('cat.jpg')

如何部署在网页呢

可以先使用自己的电脑作为服务器
首先使用pip install gradio安装好这个库
编写app.py文件在当前目录下

from fastai.vision.all import *
import gradio as gr

learn = load_learner('catorchicken.pkl')

categories = ('cat','chicken')

def classify_image(img):
    pred,idx,probs = learn.predict(img)
    return dict(zip(categories, map(float,probs)))

image = gr.Image(height=512,width=512)
label = gr.Label()
examples = ['cat1.jpg','chicken.jpg','funny.jpg']

intf = gr.Interface(fn=classify_image, inputs=image, outputs=label, examples=examples)
intf.launch(inline=False)