前言

在计算机视觉领域，图像分类是一个基础且重要的任务。本文将介绍如何使用MobileNetV3预训练模型来训练一个水果分类模型，并通过Flask框架进行部署。MobileNetV3作为轻量级网络，在保持较高精度的同时，具有较快的推理速度，非常适合实际应用场景。

环境准备

首先，我们需要准备以下环境：

# 主要依赖包
torch>=1.7.0
torchvision>=0.8.0
flask>=2.0.0
pillow>=8.0.0
numpy>=1.19.0
requests>=2.25.0  # 用于数据采集
matplotlib>=3.3.0  # 用于绘制训练曲线

数据集准备

1. 数据采集

我们使用百度图片API来采集水果图片数据。以下是数据采集的代码实现：

import requests
import osdef get_images(keyword, page_num):headers = {'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; WOW64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/69.0.3497.81 Safari/537.36'}url = 'https://image.baidu.com/search/acjson?'# 设置图片保存路径download_path = os.path.join("./data", keyword)if not os.path.exists(download_path):os.makedirs(download_path)# 构造请求参数params = {'tn': 'resultjson_com','word': keyword,'pn': 0,'rn': 30,# ... 其他参数}# 下载图片for i in range(page_num):params["pn"] = i*30response = requests.get(url, params=params, headers=headers)# 处理返回结果并保存图片

2. 数据集组织

将采集到的图片按照以下结构组织：

data/├── apple/│   ├── 0.jpg│   ├── 1.jpg│   └── ...├── banana/│   ├── 0.jpg│   ├── 1.jpg│   └── ...└── ...

模型训练

1. 数据加载和预处理

from torchvision import transforms
from torchvision.datasets import ImageFolder
from torch.utils.data import random_split, DataLoader# 图像预处理
transform = transforms.Compose([transforms.Resize(256),transforms.CenterCrop(224),transforms.ToTensor(),transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])# 加载数据集
dataset = ImageFolder("data", transform=transform)# 保存类别标签
with open("label.txt", "w", encoding="UTF-8") as f:for line in dataset.classes:f.write(line + "\n")# 划分训练集和测试集
train_ratio = 0.8
train_size = int(len(dataset) * train_ratio)
test_size = len(dataset) - train_size
train_dataset, test_dataset = random_split(dataset, [train_size, test_size])# 创建数据加载器
batch_size = 64
train_loader = DataLoader(train_dataset, batch_size=batch_size)
test_loader = DataLoader(test_dataset, batch_size=batch_size)

2. 模型定义

from torchvision import models
import torch.nn as nn# 使用MobileNetV3-Small预训练模型
model = models.mobilenet_v3_small(pretrained=True)
# 修改最后的分类层
model.classifier[3] = nn.Linear(in_features=1024, out_features=5)  # 5个类别# 如果有GPU则使用GPU
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = model.to(device)

3. 训练过程

# 定义损失函数和优化器
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)# 训练参数
num_epochs = 20
best_valid_acc = 0
best_model = None# 记录训练过程
train_losses = []
valid_losses = []
train_accs = []
valid_accs = []for epoch in range(num_epochs):# 训练阶段model.train()train_loss = 0.0train_acc = 0.0total = 0for inputs, labels in train_loader:inputs, labels = inputs.to(device), labels.to(device)optimizer.zero_grad()outputs = model(inputs)loss = criterion(outputs, labels)loss.backward()optimizer.step()train_loss += loss.item()_, predicted = torch.max(outputs, 1)train_acc += (predicted == labels).sum().item()total += len(labels)

Flask部署

1. 创建Flask应用

2. 实现预测接口

@app.route('/predict', methods=['POST'])
def predict():if 'image' not in request.files:return render_template('index.html', prediction=None)image_file = request.files['image']image_data = image_file.read()# 图像预处理img = Image.open(io.BytesIO(image_data))img = transform(img)img = torch.unsqueeze(img, dim=0)# 模型预测with torch.no_grad():prediction = model(img)prediction = F.softmax(prediction, dim=1)# 获取预测结果pred_label = class_labels[torch.argmax(prediction).item()]confidence = torch.max(prediction).item()return render_template('index.html', prediction=pred_label,confidence=confidence)