深度学习视觉全通关:从 CNN 原理到产品化部署
猫狗图片分类实战:从零手写 CNN
在进入代码实现之前,我们需要理解两个核心概念:图像分类任务本身是什么,以及 CNN(卷积神经网络) 是如何像人类大脑一样“看”世界的。
1. 图像分类概念概述
图像分类(Image Classification)是计算机视觉中最基础的任务。它的目标是将输入的图像分配给一个特定的标签(类别)。
- 输入:一张图片的像素矩阵(对于彩色图,通常是 RGB 三通道)。
- 输出:每个类别的概率得分(例如:猫 90%,狗 10%)。
- 挑战:同一个物体的光照、拍摄角度、背景遮挡都会改变像素值,但类别本质不变。
2. 卷积神经网络 (CNN) 技术简介
传统的神经网络(全连接网络)在处理图像时会丢失空间信息(把二维图片拉成一维向量)。而 CNN 通过模拟生物视觉机制,能够有效地提取图像的局部特征。
2.1 CNN 的核心组件
- 卷积层 (Convolutional Layer):
- 原理:使用一组可学习的“滤镜”(Filter)在图片上滑动。
- 作用:提取特征。底层的卷积层提取线条、边缘;深层的卷积层提取眼睛、耳朵等复杂形状。
- 激活层 (Activation Layer - 如 ReLU):
- 作用:引入非线性。它告诉模型哪些特征是“重要的”(激活),哪些是“噪音”(抑制)。
- 池化层 (Pooling Layer):
- 作用:降维(压缩)。在保留核心特征的同时,减小数据量,并赋予模型“平移不变性”(即猫在图片左边或右边都能识别出来)。
- 全连接层 (Fully Connected Layer):
- 作用:分类器。将前面提取到的所有特征汇总,最终决定这张图到底属于哪一类。
3. 为什么选择 PyTorch 手写实现?
虽然现成的模型(如 ResNet, EfficientNet)精度极高,但从零手写一个简单的 CNN 具有以下教育价值:
- 理解维度变化:亲自计算图片经过每一层后的尺寸变化。
- 掌握数据流:理解张量(Tensor)如何在模型中传递。
- 轻量化:相比动辄几百 MB 的大型模型,手写 CNN 只有几 MB,适合在普通笔记本电脑甚至手机端运行。
4. 环境准备 (2026 年推荐)
建议使用 Python 3.10+ 和最新版的 PyTorch。
- 深度学习库:
torch, torchvision
- 数据处理:
Pillow, numpy
- 可视化:
matplotlib
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121
5. 核心代码实现
本节将分为三个模块:数据增强、模型定义、以及训练主循环。
5.1 数据增强与公开数据集加载
我们使用 Modelscope的小数据集。这个数据集只有10MB,345张(275训练+70测试)
链接:https://pan.baidu.com/s/1qYa13SxFM0AirzDyFMy0mQ
提取码:1ybm [citation:2]
python:train_data_prep.py
import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms
from torch.utils.data import DataLoader
# 设置设备
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# 定义数据增强
train_transform = transforms.Compose([
transforms.Resize((128, 128)),
transforms.RandomHorizontalFlip(),
transforms.RandomRotation(15),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])
# 加载数据集
train_data = datasets.ImageFolder('data/dataset_cats_and_dogs/train', transform=train_transform)
val_data = datasets.ImageFolder('data/dataset_cats_and_dogs/val', transform=train_transform)
train_loader = DataLoader(train_data, batch_size=8, shuffle=True)
print(f"训练集大小: {len(train_data)}")
print(f"测试集大小: {len(val_data)}")
5.2 构建自定义 CNN 模型
我们将手动定义一个包含两层卷积和两层全连接的网络,加入 Dropout 来增强泛化能力。。
python:train_data_prep.py
class SimpleCNN(nn.Module):
def __init__(self):
super(SimpleCNN, self).__init__()
# 卷积层 1: 输入 3 通道 (RGB),输出 16 通道
self.conv1 = nn.Conv2d(3, 16, kernel_size=3, padding=1)
# 卷积层 2: 输入 16 通道,输出 32 通道
self.conv2 = nn.Conv2d(16, 32, kernel_size=3, padding=1)
# 最大池化层: 2x2 窗口,步长为 2
self.pool = nn.MaxPool2d(2, 2)
# 维度推导: 128 -> (pool1) 64 -> (pool2) 32
# 展平后的维度: 32通道 * 32宽 * 32高
self.fc1 = nn.Linear(32 * 32 * 32, 128)
self.dropout = nn.Dropout(0.5) # 防止过拟合
self.fc2 = nn.Linear(128, 2) # 最终输出分类:猫(0) 或 狗(1)
def forward(self, x):
x = self.pool(torch.relu(self.conv1(x)))
x = self.pool(torch.relu(self.conv2(x)))
x = x.view(-1, 32 * 32 * 32) # 展平
x = torch.relu(self.fc1(x))
x = self.dropout(x)
x = self.fc2(x)
return x
model = SimpleCNN().to(device)
5.3 训练逻辑 (Training Loop)
训练过程包括前向传播、计算损失、反向传播和权重更新,我们在训练循环中加入简单的准确率(Accuracy)计算,直观观察效果。
python:train_data_prep.py
# 定义损失函数和优化器
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)
def train_and_evaluate(epochs=10, save_path="simple_cnn_cat_dog.pth"):
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = SimpleCNN().to(device)
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
best_acc = 0.0 # 用于追踪最高准确率
for epoch in range(epochs):
model.train()
running_loss = 0.0
correct = 0
total = 0
for images, labels in train_loader:
images, labels = images.to(device), labels.to(device)
# 前向传播
outputs = model(images)
loss = criterion(outputs, labels)
# 反向传播
optimizer.zero_grad()
loss.backward()
optimizer.step()
# 统计
running_loss += loss.item()
_, predicted = torch.max(outputs.data, 1)
total += labels.size(0)
correct += (predicted == labels).sum().item()
# 计算当前 Epoch 的平均准确率
epoch_acc = 100 * correct / total
print(f"Epoch [{epoch+1}/{epochs}], Loss: {running_loss/len(train_loader):.4f}, Acc: {epoch_acc:.2f}%")
# --- 保存逻辑 ---
if epoch_acc > best_acc:
best_acc = epoch_acc
# 保存整个模型或仅保存状态字典(推荐保存 state_dict)
torch.save(model.state_dict(), save_path)
print(f"--> 检测到更好的模型,已保存至 {save_path} (Acc: {best_acc:.2f}%)")
# ----------------
print(f"训练完成!最高准确率: {best_acc:.2f}%")
train_and_evaluate()
6. 进阶技巧与优化
- Dropout(随机失活):在全连接层中加入
nn.Dropout(0.5)。这就像是考试时随机盖住一部分笔记,强迫大脑(神经元)独立思考,防止死记硬背(过拟合)。
- 批标准化 (Batch Normalization):在卷积层后加入
nn.BatchNorm2d。它能让训练过程更稳定,允许使用更高的学习率。
- 学习率调度 (Scheduler):随着训练进行逐渐减小学习率,帮助模型在后期更精准地找到“最优解”。
通过本教程,你已经实现了一个具备基础视觉能力的 CNN:
- 输入端:通过
transforms 规范化了数据。
- 特征端:通过
Conv2d 提取了空间特征。
- 决策端:通过
Linear 层完成了逻辑分类。
7. 本地图片测试脚本
在运行此脚本前,请确保你已经保存了模型权重文件(如 simple_cnn_cat_dog.pth),并且你的电脑里有一张准备测试的图片。
my_pet.jpg
7.1 推理流程概述
推理(Inference)与训练不同,它不需要计算梯度。其核心步骤如下:
- 加载模型结构:必须与训练时的网络结构完全一致。
- 加载权重:将训练好的参数填入模型。
- 图像预处理:必须使用与训练时相同的缩放和归一化参数。
- 前向传播:获取得分最高的类别。
7.2 完整测试代码
import torch
import torch.nn as nn
import torch.nn.functional as F
from torchvision import transforms
from PIL import Image
# 1. 必须定义与训练时完全一样的模型类
class SimpleCNN(nn.Module):
def __init__(self):
super(SimpleCNN, self).__init__()
self.conv1 = nn.Conv2d(3, 16, 3, padding=1)
self.conv2 = nn.Conv2d(16, 32, 3, padding=1)
self.pool = nn.MaxPool2d(2, 2)
self.fc1 = nn.Linear(32 * 32 * 32, 128)
self.dropout = nn.Dropout(0.5)
self.fc2 = nn.Linear(128, 2)
def forward(self, x):
x = self.pool(F.relu(self.conv1(x)))
x = self.pool(F.relu(self.conv2(x)))
x = x.view(-1, 32 * 32 * 32)
x = F.relu(self.fc1(x))
x = self.dropout(x)
x = self.fc2(x)
return x
def predict_local_image(image_path, model_path='simple_cnn_cat_dog.pth'):
# A. 准备设备和模型
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = SimpleCNN().to(device)
# B. 加载训练好的权重
try:
model.load_state_dict(torch.load(model_path, map_location=device))
model.eval() # 切换到评估模式(关闭 Dropout)
print("模型权重加载成功!")
except FileNotFoundError:
print(f"错误:找不到模型文件 {model_path},请先运行训练脚本。")
return
# C. 定义与训练一致的预处理
transform = transforms.Compose([
transforms.Resize((128, 128)),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])
# D. 加载并转换图片
try:
img = Image.open(image_path).convert('RGB')
img_tensor = transform(img).unsqueeze(0).to(device) # 增加 batch 维度 [1, 3, 128, 128]
except Exception as e:
print(f"图片读取失败: {e}")
return
# E. 执行推理
with torch.no_grad(): # 推理不需要计算梯度
outputs = model(img_tensor)
probabilities = F.softmax(outputs, dim=1) # 转化为概率
confidence, predicted = torch.max(probabilities, 1)
# F. 输出结果
classes = ['猫 (Cat)', '狗 (Dog)']
result = classes[predicted.item()]
print("-" * 30)
print(f"预测结果: {result}")
print(f"置信度: {confidence.item() * 100:.2f}%")
print("-" * 30)
# 使用示例:将 'test.jpg' 换成你电脑里的图片路径
if __name__ == "__main__":
predict_local_image('my_pet.jpg')
7.3 常见问题排查
- 结果全是“猫”或全是“狗”:这通常是因为训练轮数(Epochs)太少,或者数据集里的图片不够多样。手写小模型需要较多数据才能学会泛化。
- 尺寸报错:如果提示
size mismatch,请检查 predict 脚本中的 transforms.Resize 是否与你训练时的设置完全一致(教程中我们改成了 128x128)。
- 模型文件后缀:PyTorch 习惯使用
.pth 或 .pt 作为权重后缀,两者没有本质区别。
下一步建议:
如果你想让准确率从目前的水平(可能只有 60%-70%)提升到 95% 以上,我们可以聊聊如何将这个推理脚本应用到 ResNet18 迁移学习 模型上。你想试试看“降维打击”的威力吗?
既然你已经准备好见识“工业级”的力量,我们将模型从零手写的 SimpleCNN 升级为 ResNet18。
ResNet(残差网络)通过“跳跃连接”解决了深层网络难以训练的问题。使用预训练模型就像是聘请了一位已经读完百科全书(ImageNet 数据集)的博士,你只需要教他如何分辨特定的“猫”和“狗”品种即可。
8. 迁移学习 (Transfer Learning) 实战
8.1 为什么 ResNet18 效果更好?
- 特征提取能力:它已经在包含 1000 个类别的 120 万张图片上训练过,它对边缘、纹理、眼睛、耳朵等特征极其敏感。
- 残差结构:它允许神经网络变得更深而不会产生梯度消失问题。
8.2 ResNet18 完整训练脚本
这段代码演示了如何“改造” ResNet18 的输出层并进行微调(Fine-tuning)。
在下面的代码中,weights=models.ResNet18_Weights.DEFAULT 是关键。
冻结 vs 微调:
冻结:如果你只想训练最后一层(分类头),可以设置 for param in model.parameters(): param.requires_grad = False,然后只更新 model.fc。
微调(我们的做法):我们让整个网络以较小的学习率(0.0001)重新学习,这能让模型在保持原有视觉能力的同时,更契合你提供的猫狗数据集。
import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import models
from torchvision import datasets, transforms
from torch.utils.data import DataLoader
# 设置设备
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# 定义数据增强
train_transform = transforms.Compose([
transforms.Resize((128, 128)),
transforms.RandomHorizontalFlip(),
transforms.RandomRotation(15),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])
# 加载数据集
train_data = datasets.ImageFolder('data/dataset_cats_and_dogs/train', transform=train_transform)
val_data = datasets.ImageFolder('data/dataset_cats_and_dogs/val', transform=train_transform)
print(f"训练集大小: {len(train_data)}")
print(f"测试集大小: {len(val_data)}")
train_loader = DataLoader(train_data, batch_size=8, shuffle=True)
# 1. 加载模型:使用预训练的 ResNet18 权重
# 这就像给模型装上了一个已经训练好的“大脑”
model = models.resnet18(weights=models.ResNet18_Weights.DEFAULT)
# 2. 修改最后的全连接层
# 原本 ResNet18 输出 1000 类,我们需要将其改为 2 类(猫和狗)
num_ftrs = model.fc.in_features
model.fc = nn.Linear(num_ftrs, 2)
model = model.to(device)
# 3. 定义损失函数和优化器
criterion = nn.CrossEntropyLoss()
# 学习率通常设得比手写 CNN 小一些(如 0.0001),因为模型已经是半成品了
optimizer = optim.Adam(model.parameters(), lr=0.0001)
# 4. 训练函数
def train_model(epochs=5):
for epoch in range(epochs):
model.train()
running_loss = 0.0
correct = 0
total = 0
for images, labels in train_loader:
images, labels = images.to(device), labels.to(device)
# 梯度清零
optimizer.zero_grad()
# 前向传播
outputs = model(images)
loss = criterion(outputs, labels)
# 反向传播与优化
loss.backward()
optimizer.step()
running_loss += loss.item()
_, predicted = torch.max(outputs.data, 1)
total += labels.size(0)
correct += (predicted == labels).sum().item()
print(
f"Epoch [{epoch + 1}/{epochs}] - Loss: {running_loss / len(train_loader):.4f} - Acc: {100 * correct / total:.2f}%")
# 5. 保存训练好的权重,供 Gradio 或推理脚本使用
torch.save(model.state_dict(), 'resnet18_cat_dog.pth')
print("ResNet18 模型已保存为 resnet18_cat_dog.pth")
if __name__ == '__main__':
train_model(epochs=5)
8.3 修改后的推理脚本
我们将之前的 SimpleCNN 替换为官方预训练的 resnet18。注意:代码量反而变少了。
import torch
import torch.nn as nn
from torchvision import models, transforms
from PIL import Image
def predict_with_resnet(image_path, model_path='resnet18_cat_dog.pth'):
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# 1. 加载预训练模型结构
# weights=None 表示我们只想要结构,稍后手动加载自己的权重
model = models.resnet18(weights=None)
# 2. 修改最后的全连接层 (必须与训练时一致)
num_ftrs = model.fc.in_features
model.fc = nn.Linear(num_ftrs, 2)
# 3. 加载权重
model.load_state_dict(torch.load(model_path, map_location=device))
model = model.to(device)
model.eval()
# 4. 预处理 (ResNet 官方推荐参数)
transform = transforms.Compose([
transforms.Resize(256),
transforms.CenterCrop(224), # ResNet 通常输入 224x224
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])
# 5. 读取与预测
img = Image.open(image_path).convert('RGB')
img_tensor = transform(img).unsqueeze(0).to(device)
with torch.no_grad():
outputs = model(img_tensor)
_, predicted = torch.max(outputs, 1)
classes = ['猫 (Cat)', '狗 (Dog)']
print(f"ResNet18 预测结果: {classes[predicted.item()]}")
predict_with_resnet('my_pet.jpg')
9. 性能对比表
如果说 CNN 是通过“局部观察”来识图,那么 ViT (Vision Transformer) 就是通过“全局注意力”来识图。它是目前计算机视觉(CV)领域的 SOTA(State-of-the-art)架构,性能甚至超越了传统的 ResNet。
10.1 ViT 的核心概念:图片即单词
在 ViT 出现之前,Transformer 主要用于处理文字(如 ChatGPT)。ViT 的奇思妙想在于:
- 图像切片 (Patching):将一张图片切成 16×16 个小方块。
- 线性投影 (Embedding):把每个小方块看作一个“单词”。
- 自注意力机制 (Self-Attention):让每个小方块去观察其他所有方块,从而理解图片的全局结构(例如:左上角的猫耳和右下角的猫尾是如何关联的)。
10.2 实现代码:ViT 迁移学习
在 PyTorch 中,使用 ViT 同样非常简洁。我们使用 vit_b_16(Base 版本,16x16 切片大小)。
A. 修改模型加载部分
import torch
import torch.nn as nn
from torchvision import models
def get_vit_model(num_classes=2):
# 1. 加载预训练的 ViT Base 模型 (基于 ImageNet 21k 预训练)
weights = models.ViT_B_16_Weights.DEFAULT
model = models.vit_b_16(weights=weights)
# 2. 修改分类头 (Heads)
# ViT 的最后输出层叫做 'heads',是一个序列容器
in_features = model.heads.head.in_features
model.heads.head = nn.Linear(in_features, num_classes)
return model
# model = get_vit_model().to(device)
B. 注意事项:数据预处理
1. 核心训练逻辑
ViT 模型通常需要更小的学习率。如果学习率太大,Transformer 的自注意力层(Self-Attention)很容易导致梯度爆炸或模型不收敛。
2. ViT 训练的“闭坑”指南
尺寸匹配: ViT 内部将图片切成 16×16 的方块。如果输入不是 224 的倍数(比如你之前的 128),模型虽然能跑,但效果会大打折扣。请务必在 transforms 中使用 224x224。
显存开销: ViT 极其吃显存。如果你的显存小于 8GB,且出现 Out of Memory (OOM) 报错,请尝试将 batch_size 减小到 4 甚至 2。
收敛速度: Transformer 架构没有 CNN 那样的“感官野”(Receptive Field)先验知识,因此在前几轮训练中,Loss 下降可能比 ResNet 慢,但一旦收敛,它的上限通常更高。
3. ViT 对输入图片的尺寸和归一化非常挑剔。通常需要 224x224 或 384x384 的固定尺寸。
# ViT 专用预处理
vit_transforms = transforms.Compose([
transforms.Resize((224, 224)),
transforms.ToTensor(),
transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]) # ViT 常用归一化
])
C. 完整训练代码
import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import models, transforms, datasets
from torch.utils.data import DataLoader
# 1. 适配 ViT 的预处理 (注意:ViT 官方推荐尺寸通常是 224x224)
vit_transform = transforms.Compose([
transforms.Resize((224, 224)),
transforms.ToTensor(),
transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5]) # ViT 习惯的归一化
])
# 设置设备
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# 重新加载你的数据集(使用新尺寸)
train_data = datasets.ImageFolder('data/dataset_cats_and_dogs/train', transform=vit_transform)
train_loader = DataLoader(train_data, batch_size=8, shuffle=True)
# 2. 初始化预训练 ViT 模型
def create_vit_model():
# 加载预训练权重
model = models.vit_b_16(weights=models.ViT_B_16_Weights.DEFAULT)
# 修改分类头 (Heads)
# ViT 的输出层是一个包含 head 的序列模型
num_in_features = model.heads.head.in_features
model.heads.head = nn.Linear(num_in_features, 2) # 猫狗 2 分类
return model
model = create_vit_model().to(device)
# 3. 损失函数与优化器
criterion = nn.CrossEntropyLoss()
# 关键点:对于 ViT,学习率通常设置得非常低(如 1e-5 或 5e-5)
optimizer = optim.Adam(model.parameters(), lr=1e-5)
# 4. 训练函数
def train_vit(epochs=5):
print("开始训练 ViT 模型...")
for epoch in range(epochs):
model.train()
running_loss = 0.0
correct = 0
total = 0
for i, (images, labels) in enumerate(train_loader):
images, labels = images.to(device), labels.to(device)
optimizer.zero_grad()
outputs = model(images)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
running_loss += loss.item()
_, predicted = torch.max(outputs.data, 1)
total += labels.size(0)
correct += (predicted == labels).sum().item()
if (i + 1) % 5 == 0:
print(f"Batch [{i + 1}/{len(train_loader)}] Loss: {loss.item():.4f}")
epoch_acc = 100 * correct / total
print(
f"Epoch [{epoch + 1}/{epochs}] 完成! 平均 Loss: {running_loss / len(train_loader):.4f}, 准确率: {epoch_acc:.2f}%")
# 保存模型
torch.save(model.state_dict(), 'vit_cat_dog.pth')
print("ViT 模型权重已成功保存。")
if __name__ == '__main__':
train_vit(epochs=5)
10.3 ViT 推理脚本 (本地测试)
你可以直接运行这个脚本来感受 ViT 的识别精度:
import torch
from torchvision import models, transforms
from PIL import Image
def predict_with_vit(image_path, model_path='vit_cat_dog.pth'):
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# 初始化模型结构
model = models.vit_b_16(weights=None)
model.heads.head = torch.nn.Linear(model.heads.head.in_features, 2)
# 加载你的训练权重 (假设你已经训练并保存了)
model.load_state_dict(torch.load(model_path, map_location=device))
model.to(device).eval()
# 预处理
transform = transforms.Compose([
transforms.Resize((224, 224)),
transforms.ToTensor(),
transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])
])
img = Image.open(image_path).convert('RGB')
img_tensor = transform(img).unsqueeze(0).to(device)
with torch.no_grad():
output = model(img_tensor)
prob = torch.softmax(output, dim=1)
score, pred = torch.max(prob, 1)
classes = ['猫 (Cat)', '狗 (Dog)']
print(f"ViT 预测: {classes[pred.item()]} (置信度: {score.item() * 100:.2f}%)")
predict_with_vit('my_pet.jpg')
10.4 CNN vs ViT:你应该选哪个?
11. 总结
从 SimpleCNN 到 ResNet 再到 ViT,你已经完成了从“手工作坊”到“工业流水线”再到“未来科技”的跨越:
- SimpleCNN 让你理解了特征提取的基本原理。
- ResNet 展示了通过残差连接(Skip Connections)构建深层网络的威力。
- ViT 证明了全局注意力机制可以彻底改变我们处理图像的方式。
下一步建议:
由于 ViT 的训练非常吃配置,如果你觉得本地跑不动,我们可以研究如何使用 Gradio 为你的模型做一个可视化 Web 界面,让你能在手机浏览器上上传图片并查看预测结果。你想试试看吗?
极简 Web 应用:使用 Gradio 部署你的猫狗分类器 description: 为 PyTorch 模型构建可视化交互界面(2026 年最新版)
11. 将模型转化为 Web 应用 (Gradio)
在完成 SimpleCNN、ResNet 或 ViT 的训练后,最令人兴奋的一步就是让非技术用户也能使用它。Gradio 是目前最流行的机器学习演示框架,它可以让你用几行 Python 代码就生成一个漂亮的网页界面。
11.1 为什么使用 Gradio?
- 无需前端知识:不需要写 HTML/CSS/JS。
- 自动生成外网链接:通过
share=True 参数,你可以即时生成一个有效期 72 小时的外网访问链接,发给手机或其他设备测试。
- 内置组件:自带图片上传框、进度条、标签输出等。
11.2 环境准备
你需要安装 Gradio 库:
11.3 完整实现代码
我们将使用之前的 ViT 或 ResNet 模型作为逻辑核心。
import torch
import torch.nn as nn
import torch.nn.functional as F
from torchvision import models, transforms
import gradio as gr
from PIL import Image
# 1. 重新配置模型 (以 ResNet18 为例)
def load_trained_model(model_path='resnet18_cat_dog.pth'):
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = models.resnet18(weights=None)
model.fc = nn.Linear(model.fc.in_features, 2)
# 加载权重
model.load_state_dict(torch.load(model_path, map_location=device))
model.to(device).eval()
return model, device
model, device = load_trained_model()
# 2. 定义预处理
transform = transforms.Compose([
transforms.Resize((224, 224)),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])
# 3. 定义预测函数 (供 Gradio 调用)
def classify_image(inp_img):
if inp_img is None:
return None
# 图片预处理
inp_img = Image.fromarray(inp_img) # Gradio 传入的是 numpy 数组
img_tensor = transform(inp_img).unsqueeze(0).to(device)
# 推理
with torch.no_grad():
output = model(img_tensor)
prediction = F.softmax(output, dim=1)[0] # 转化为概率
# 构建返回结果:{类别: 概率}
classes = ['猫 (Cat)', '狗 (Dog)']
return {classes[i]: float(prediction[i]) for i in range(2)}
# 4. 构建 Gradio 界面
demo = gr.Interface(
fn=classify_image, # 点击提交后调用的函数
inputs=gr.Image(), # 输入组件:图片上传
outputs=gr.Label(num_top_classes=2), # 输出组件:概率标签
title="🐱 智能猫狗识别系统 🐶",
description="上传一张猫或狗的照片,AI 会告诉你它的类别及置信度。",
examples=["cat_example.jpg", "dog_example.jpg"] # 可选:添加示例图片
)
# 5. 启动应用
if __name__ == "__main__":
# share=True 会生成一个公网访问链接
demo.launch(share=True)
11.4 交互功能解析
- 输入端 (
gr.Image):支持拖拽图片、点击上传,甚至在手机端直接调用摄像头拍摄。
- 输出端 (
gr.Label):会自动显示出概率最高的类别,并以进度条形式展示置信度。
- 共享模式 (
share=True):Gradio 会建立一个 SSH 隧道。当你运行代码后,控制台会输出一个类似 https://xxxx.gradio.live 的链接。只要你的电脑不关,世界上任何人都能通过这个链接访问你的模型。
12. 整个系列总结:你的 AI 进阶之路
回顾这一系列教程,你已经从最基础的卷积神经网络进化到了顶尖的架构:
- 初级:手写
SimpleCNN —— 掌握像素到特征的数学转换。
- 中级:迁移学习
ResNet18 —— 学习如何站在巨人的肩膀上解决实际问题。
- 高级:
Vision Transformer —— 探索基于注意力机制的最新 CV 范式。
- 实战:
Gradio 部署 —— 将纯代码转化为可交付的产品。
恭喜!你已经完成了一个从算法研发到产品原型发布的完整全栈深度学习流程。