大一小白从零到一:我用ResNet18做了一个猫狗识别AI
大家好,我是一个普通一本的大一计科生。这篇文章不是大佬的教学,而是一个小白踩坑一个月的真实记录。一个月前我连 import torch 都会报错,现在做出了一个能上传图片判断猫狗的网站。

如果你也刚入坑AI,想做一个完整的项目但不知道从哪下手,这篇文章应该能帮到你。
项目效果

一、项目整体思路
这个项目分三步:数据 → 模型 → 网站
爬虫下载猫狗图片 → 训练ResNet18模型 → Flask搭建网站 → 上传图片识别
技术栈就是PyTorch加Flask,没有用很花哨的东西。
技术栈一览
爬虫 requests + re 轻量,不用Scrapy那么重
训练 PyTorch + ResNet18 预训练模型,省时间
网站 Flask + HTML 简单,几行代码就能跑
部署 本地运行 先跑通再说
二、第一步:爬虫采集数据
为啥不直接用现成的数据集?
网上有猫狗数据集,但我想自己控制数据来源。而且爬虫写好了,以后想加"兔子""鸟"之类的,改个关键词就能跑。
代码实现
python
import requests
import re
import os
import time
def download_baidu_images(keyword, save_dir, num=30):
os.makedirs(save_dir, exist_ok=True)
headers = {
'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36'
}
downloaded = 0
page = 0
while downloaded < num:
url = f'https://image.baidu.com/search/flip?tn=baiduimage&word={keyword}&pn={page*20}'
resp = requests.get(url, headers=headers)
# 从返回的html里抠图片链接
img_urls = re.findall(r'"objURL":"(.*?)"', resp.text)
for img_url in img_urls:
if downloaded >= num:
break
try:
img_data = requests.get(img_url, timeout=10)
with open(f'{save_dir}/{keyword}_{downloaded+1}.jpg', 'wb') as f:
f.write(img_data.content)
downloaded += 1
print(f'已下载 {downloaded}/{num}')
time.sleep(0.5) # 别太快,会被封
except:
continue
page += 1
time.sleep(1)
download_baidu_images('猫', './dataset/猫', 30)
download_baidu_images('狗', './dataset/狗', 30)
踩过的坑
1. 百度会封IP:加 time.sleep 和 User-Agent 会好很多。封了等几分钟再试。
2. 有些链接不是图片:下载下来打不开的,直接 except 跳过就行。
3. 图片格式不对:统一保存成 .jpg,PIL都能打开。
三、第二步:训练模型
为啥选ResNet18?
ResNet18够用,ResNet50我的电脑带不动。而且PyTorch一行代码就能下载预训练好的版本,直接改最后一层就行。
数据预处理
python
from torchvision import transforms
transform = transforms.Compose([
transforms.Resize((224, 224)), # ResNet输入要求
transforms.RandomHorizontalFlip(), # 随机翻转,相当于白嫖数据
transforms.ToTensor(),
transforms.Normalize([0.485,0.456,0.406], [0.229,0.224,0.225]) # ImageNet标准
])
模型定义
python
import torch.nn as nn
import torchvision.models as models
model = models.resnet18(pretrained=True) # 下载预训练权重
model.fc = nn.Linear(512, 2) # 改成2分类:猫和狗
训练代码
python
import torch.optim as optim
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)
for epoch in range(10):
model.train()
total_loss = 0
for images, labels in train_loader:
optimizer.zero_grad()
outputs = model(images)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
total_loss += loss.item()
print(f'Epoch {epoch+1}, Loss: {total_loss/len(train_loader):.4f}')
训练结果
我只有60张图(猫30,狗30),训练了10个epoch:
· 训练准确率:100%(已经背下来了)
· 验证准确率:83.33%(12张验证集猜对了10张)
数据量太小,有过拟合。如果你有200张以上,准确率应该能到90%以上。
四、第三步:部署成网站
为啥用Flask?
因为简单。整个后端就几十行代码,不需要学Django那种大家伙。
后端代码
python
from flask import Flask, request, render_template_string
import torch
from PIL import Image
app = Flask(__name__)
model = load_model() # 启动时加载一次就行
@app.route('/')
def index():
return render_template_string(HTML) # 前端代码直接写在里面
@app.route('/predict', methods=['POST'])
def predict():
file = request.files['file']
img = Image.open(file.stream).convert('RGB')
# 预处理 → 预测 → 返回结果
return {'class': '猫', 'confidence': 0.99}
```
前端界面
就是一个HTML页面:上传框 + 按钮 + 结果显示区。支持点击上传和拖拽上传。
前端代码比较长,核心就是用 FormData 发请求:
```javascript
fetch('/predict', {
method: 'POST',
body: formData
})
.then(res => res.json())
.then(data => {
document.getElementById('result').innerHTML =
`识别结果:${data.class},置信度:${(data.confidence*100).toFixed(2)}%`;
});
```
启动
```bash
python app.py
```
浏览器打开 http://127.0.0.1:5000,上传一张图,就能看到结果了。
五、项目优化方向
目前只是个demo,想做得更好可以从这几个方向下手:
1. 增加数据量:爬虫多跑几次,每类100张以上,准确率能到90%以上
2. 加更多类别:爬虫关键词改成"鸟""兔子",重新训练就行
3. 部署到云服务器:让别人可以点链接直接访问,不用自己跑代码
六、一些真心话
一个月能做成这样?
能。前提是每天坚持2-3小时。我大一课不多,晚上和周末都在写代码。
最开始也看不懂 view(-1, 784) 是啥意思,squeeze() 为啥要去掉维度1。硬着头皮跑通,跑多了就懂了。
给同样在入门的同学几点建议
1. 不要一直看课:看十节课不如跑通一个项目
2. 代码一定要手敲:复制粘贴看起来快,但坑都踩不到,等于白学
3. 遇到报错读英文:别直接复制到百度,试着看最后一行错误信息,慢慢就能看懂
4. 做完一定要写文章:写下来才是自己的。我在写这篇文章的时候,又把代码理了一遍,发现好几个之前没注意的点
关于大一找实习
有人问我现在能不能去实习。实话实说:挺难的。不是因为能力不够,是企业根本不招大一。
但这不代表暑假没事干。我计划暑假:刷LeetCode + 做第二个项目(试试YOLO目标检测)+ 写技术博客。大二暑假再去找实习。
附:环境配置
# 创建虚拟环境
conda create -n ai_env python=3.9
conda activate ai_env
# 安装依赖
pip install torch torchvision flask pillow requests tqdm
# 运行
python spider.py --keyword "猫" --num 30 # 爬数据
python train.py # 训练
python app.py # 启动网站
`
更多推荐


所有评论(0)