ARTICLE · INTELLIGENCE

战地情报 · 详情页

来自尧图项目组的一线实战观察与深度解析

PyTorch垃圾分类模型本地部署实战:Flask+Docker一键运行

PyTorch垃圾分类模型本地部署实战:Flask+Docker一键运行 简介本资源是一份面向高校Python课程学习者的深度学习实践项目聚焦垃圾分类这一典型计算机视觉应用场景适合具备基础Python与PyTorch/TensorFlow知识的本科生完成课程大作业或课设实践。压缩包共134个文件、75.59MB涵盖20个核心Python源码含模型训练、推理与Web接口、13个Jupyter Notebook含数据预处理、模型调优与结果可视化、12个VueJS前端页面支持图片上传与分类展示、8个Markdown文档含环境配置与模块说明以及PPT答辩稿、Word版实践报告、Dockerfile和ONNX模型等完整交付物。已有89人学习下载所有代码均经本地实测可运行通过助教审核评审得分95分以上配套部署指南清晰覆盖从环境搭建、模型加载到前后端联调的全流程目录结构分层合理便于按模块理解系统架构与工程实现逻辑。1. 垃圾分类系统不是 demo它真能跑通摄像头模型界面三件套且不用改一行代码就能在你笔记本上启动这不是一个“调用 API 打印 hello world”的 Python 小练习。这是温州大学某届课程实践里拿 95 分以上的完整闭环项目——从你用手机拍一张饮料瓶照片开始到系统弹出「可回收物塑料瓶」的红色提示框结束全程本地运行、不依赖任何云服务、不调用第三方识别接口。它用的是 ResNet18 微调模型非预训练黑盒训练数据集是公开的 TrashNet 子集 自采 200 张实拍图含反光、遮挡、角度倾斜部署方式明确支持两种路径纯 Python Flask Web 服务适合调试和 Docker 容器化适合交作业/演示。所有源码都经过助教手把手过审.gitignore里删干净了.pyc和__pycache__Dockerfile里 pip 源已切为清华镜像连requirements.txt中torch1.13.1cpu的版本号都精确锁定——不是“大概能跑”是“你解压后 cd 进目录敲python app.py就能看见 localhost:5000 的上传页”。如果你正卡在课程大作业 deadline 前三天既没时间从零搭环境又怕 PyTorch 版本冲突导致ImportError: cannot import name MultiScaleDeformableAttention这个包就是你的后悔药。2. 从源码结构到模型推理看清这 7 个文件怎么咬合在一起2.1 核心目录树别急着 run先认全这七个关键角色解压后你会看到一个扁平但逻辑清晰的结构。重点盯住以下 7 个文件其余如.editorconfig是编辑器配置semantic.css是前端样式可暂忽略文件名类型作用是否必须修改model.pyPython 模块定义 ResNet18 修改版最后一层 fc 替换为 4 分类可回收/有害/湿垃圾/干垃圾含forward()和load_state_dict()调用逻辑否已适配 TrashNet 标签train.pyPython 脚本训练入口加载data/下的 train/val 子目录设置batch_size32,epochs50,lr0.001用torch.optim.AdamCrossEntropyLoss否参数已调优若换数据集才需改app.pyPython 脚本Flask Web 入口定义/uploadPOST 接口调用predict_image()函数返回 JSON 结果并渲染templates/index.html否路径硬编码为static/uploads/但已建好目录predict.pyPython 模块推理核心含predict_image(image_path)函数完成图像读取→归一化mean[0.485,0.456,0.406], std[0.229,0.224,0.225]→模型前向→argmax 取类别否预处理参数与训练一致requirements.txt文本文件依赖清单明确列出Flask2.2.5,torch1.13.1cpu,torchvision0.14.1cpu,Pillow9.5.0,numpy1.23.5否版本锁死避免torch与torchvision不兼容Dockerfile构建脚本多阶段构建第一阶段FROM pytorch/pytorch:1.13.1-cpu安装依赖第二阶段FROM python:3.8-slim复制编译产物最终镜像仅 428MB否基础镜像已指定 CPU 版省去 CUDA 驱动适配data/目录文件夹训练数据根目录含train/4 个子文件夹每类 300 图、val/每类 100 图、test/200 张未标注图用于 demo是若要加新类别必须在此结构下新增文件夹提示data/目录是整个系统的数据心脏。它不包含原始图片而是提供了一个标准组织方式——你只需把新采集的垃圾照片按类别放入对应子文件夹如data/train/有害垃圾/废电池_001.jpgtrain.py就能自动识别并参与训练。这种结构直接复用 PyTorchImageFolder比手动写Dataset类少踩 80% 的路径 bug。2.2 模型加载与推理链为什么predict.py里model.eval()不能删predict.py的核心函数predict_image(image_path)看似简单但每一步都卡着深度学习部署的命门import torch from PIL import Image import numpy as np from model import ResNet18Classifier # 注意不是 from torchvision.models import resnet18 def predict_image(image_path): # 1. 加载图像并转为 tensor注意尺寸 img Image.open(image_path).convert(RGB) # 强制转 RGB避免 RGBA 导致通道数错误 img img.resize((224, 224), Image.BILINEAR) # 必须 resize 到 224x224ResNet 输入要求 transform transforms.Compose([ transforms.ToTensor(), # 转 tensor值域 [0,1] transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]) # 归一化到 [-1,1] 附近 ]) img_tensor transform(img).unsqueeze(0) # 增加 batch 维度(1, 3, 224, 224) # 2. 加载模型关键CPU 模式 eval 模式 model ResNet18Classifier(num_classes4) model.load_state_dict(torch.load(weights/best_model.pth, map_locationcpu)) # map_locationcpu 防止 GPU 机器报错 model.eval() # 必须否则 BatchNorm 层会用训练时的统计量Dropout 会随机失活 # 3. 推理并返回结果 with torch.no_grad(): # 关闭梯度计算节省显存/CPU output model(img_tensor) probabilities torch.nn.functional.softmax(output, dim1) pred_class torch.argmax(probabilities, dim1).item() confidence probabilities[0][pred_class].item() class_names [可回收物, 有害垃圾, 湿垃圾, 干垃圾] return { class: class_names[pred_class], confidence: round(confidence, 3), probabilities: {name: round(float(prob), 3) for name, prob in zip(class_names, probabilities[0])} }这段代码的玄学点在于model.eval()和torch.no_grad()的组合。如果删掉model.eval()BatchNorm 层会继续使用训练时保存的 running_mean/run_var而你在推理时输入单张图其 mini-batch size1会导致 BN 层输出剧烈抖动分类结果完全不可信如果删掉torch.no_grad()PyTorch 会保留计算图不仅慢 3 倍还可能因内存不足直接 crash。我第一次跑 demo 时就漏了eval()同一张塑料瓶图反复预测出「湿垃圾」「干垃圾」来回跳折腾两小时才发现是这行注释被我当成冗余删了。2.3 Flask Web 服务app.py如何把模型变成网页按钮app.py是整个系统最薄但最关键的胶水层。它不做训练、不碰模型结构只做三件事接收文件、调用预测、返回页面。它的健壮性体现在对异常的穷举覆盖from flask import Flask, request, render_template, jsonify, redirect, url_for import os from predict import predict_image from werkzeug.utils import secure_filename app Flask(__name__) UPLOAD_FOLDER static/uploads ALLOWED_EXTENSIONS {png, jpg, jpeg, bmp} # 确保上传目录存在 os.makedirs(UPLOAD_FOLDER, exist_okTrue) def allowed_file(filename): return . in filename and \ filename.rsplit(., 1)[1].lower() in ALLOWED_EXTENSIONS app.route(/) def index(): return render_template(index.html) app.route(/upload, methods[POST]) def upload_file(): if file not in request.files: return jsonify({error: No file part}), 400 file request.files[file] if file.filename : return jsonify({error: No selected file}), 400 if file and allowed_file(file.filename): # 安全重命名防止 ../ 路径遍历攻击 filename secure_filename(file.filename) filepath os.path.join(UPLOAD_FOLDER, filename) file.save(filepath) try: result predict_image(filepath) return jsonify(result) except Exception as e: # 捕获模型加载失败、图像损坏等所有异常 return jsonify({error: fPrediction failed: {str(e)}}), 500 else: return jsonify({error: File type not allowed}), 400 if __name__ __main__: app.run(host0.0.0.0, port5000, debugFalse) # debugFalse 关闭 Werkzeug 交互式 debugger防生产环境泄露这里secure_filename()是安全底线——它会把../../../etc/passwd.jpg这种恶意文件名转成etc_passwd.jpg彻底堵死路径遍历漏洞。而debugFalse是交作业前必须检查的开关开启 debug 模式会在浏览器暴露完整的 traceback助教一眼就能看到你requirements.txt里漏装了Pillow。我见过三个同学因为没关 debug答辩时被问“你这个ModuleNotFoundError是怎么产生的”当场扣分。3. Docker 部署为什么Dockerfile里要用多阶段构建而不是一行pip install3.1 Dockerfile 逐行拆解每一行都在解决一个真实痛点这个Dockerfile不是照抄教程而是针对课程作业场景定制的# 第一阶段构建环境含编译依赖 FROM pytorch/pytorch:1.13.1-cpu # 设置国内源加速 pip RUN pip config set global.index-url https://pypi.tuna.tsinghua.edu.cn/simple/ # 安装构建期依赖如 Pillow 编译需要 jpeg-dev RUN apt-get update apt-get install -y \ libjpeg-dev \ libpng-dev \ rm -rf /var/lib/apt/lists/* # 复制 requirements 并安装此时 wheel 包会编译 COPY requirements.txt . RUN pip install --no-cache-dir -r requirements.txt # 复制源码 COPY . /app WORKDIR /app # 运行测试确保模型能加载、Flask 能启动 RUN python -c import torch; print(PyTorch OK) \ python -c from app import app; print(Flask OK) \ echo Build stage passed # 第二阶段运行环境精简镜像 FROM python:3.8-slim # 复制第一阶段编译好的依赖和代码 COPY --from0 /usr/local/lib/python3.8/site-packages /usr/local/lib/python3.8/site-packages COPY --from0 /usr/local/bin/pip /usr/local/bin/pip COPY --from0 /app /app # 创建非 root 用户提升安全性课程作业虽不强制但体现工程素养 RUN useradd -m -u 1001 -G root appuser USER appuser # 暴露端口 EXPOSE 5000 # 启动命令 CMD [gunicorn, --bind, 0.0.0.0:5000, --workers, 2, app:app]关键点解析多阶段构建第一阶段用pytorch/pytorch:1.13.1-cpu含编译工具链第二阶段用python:3.8-slim仅运行时依赖最终镜像体积从 1.2GB 压到 428MB。这意味着你交作业时docker build耗时从 8 分钟降到 2 分钟助教拉取镜像也快得多。--no-cache-dir禁用 pip 缓存避免不同机器上缓存污染导致安装失败。我曾因缓存里混入了torch-2.0.0的 wheel导致import torch报undefined symbol: _ZNK3c104Type13isSubtypeOfERKS0_。gunicorn替代flask runflask run是开发服务器不支持多 worker高并发下会阻塞gunicorn开 2 个 worker能同时处理上传请求和预测请求避免用户点击上传后页面假死。3.2 构建与运行三行命令走完全流程别被 Docker 吓住课程作业级部署只需三步# 1. 构建镜像注意最后的点表示上下文是当前目录 docker build -t waste-classifier . # 2. 运行容器映射 5000 端口挂载 data 目录便于后续训练 docker run -d -p 5000:5000 -v $(pwd)/data:/app/data --name classifier-container waste-classifier # 3. 查看日志确认启动成功 docker logs classifier-container # 正常输出应含[INFO] Starting gunicorn 21.2.0 # [INFO] Listening at: http://0.0.0.0:5000注意-v $(pwd)/data:/app/data这个挂载非常关键。它让容器内app.py读取的data/目录实际指向你宿主机的data/文件夹。这样你本地新增训练图片容器内train.py就能立刻看到无需重新构建镜像。这是课程作业迭代最快的路径——改数据、跑训练、更新模型权重全程不碰 Docker。3.3 验证容器是否真在干活curl 测试比打开浏览器更可靠别急着开浏览器先用curl确认服务底层通不通# 上传一张测试图用项目自带的 test/ 下任意 jpg curl -X POST http://localhost:5000/upload \ -F filedata/test/plastic_bottle_001.jpg \ -H Content-Type: multipart/form-data预期返回{ class: 可回收物, confidence: 0.923, probabilities: { 可回收物: 0.923, 有害垃圾: 0.012, 湿垃圾: 0.031, 干垃圾: 0.034 } }如果返回{error: Prediction failed: ...}说明模型加载或图像处理出问题如果返回curl: (7) Failed to connect to localhost port 5000: Connection refused说明容器没起来或端口没映射对。这个测试绕过了浏览器缓存、JavaScript 错误等干扰项直击服务核心。我帮三个同学 debug 过其中两人都是docker run忘了-p 5000:5000curl一试就定位了。4. 避坑指南95 分作业背后的 5 个血泪经验4.1 现象ImportError: cannot import name MultiScaleDeformableAttention原因requirements.txt里torch和torchvision版本不匹配。torchvision0.14.1cpu必须搭配torch1.13.1cpu若你本地已装torch2.0.0pip install -r requirements.txt会降级失败残留旧版本。解决先彻底卸载pip uninstall torch torchvision torchaudio再执行pip install torch1.13.1cpu torchvision0.14.1cpu --index-url https://download.pytorch.org/whl/cpu。注意--index-url参数必须带否则 pip 会从 pypi.org 拉取 CUDA 版本。4.2 现象Flask 启动后访问http://localhost:5000显示404 Not Found原因app.py中render_template(index.html)要求templates/index.html存在但解压后该文件可能因压缩软件编码问题损坏尤其 Windows 下用 WinRAR 解压时。index.html里script srcstatic/js/main.js的路径若写成./static/js/main.js也会 404。解决用 VS Code 打开templates/index.html确认首行是!DOCTYPE html且无乱码检查script标签路径是否为static/js/main.js无前导./在终端执行ls -l templates/确认文件权限为-rw-r--r--。4.3 现象上传图片后返回{error: Prediction failed: expected str, bytes or os.PathLike object, not None}原因predict_image()函数中torch.load(weights/best_model.pth)找不到权重文件。weights/目录为空或best_model.pth被压缩软件过滤某些 ZIP 工具默认不打包空目录或隐藏文件。解决解压后立即执行ls -la weights/确认best_model.pth存在且大小 10MB若为空从参考报告.docx末尾的网盘链接补下载文档里写了备用下载地址若仍失败手动运行python train.py训练 5 个 epoch 生成新权重。4.4 现象Docker 容器启动后docker logs显示OSError: [Errno 99] Cannot assign requested address原因app.py中app.run(host0.0.0.0, port5000)在容器内正常但gunicorn启动命令写成了gunicorn --bind 127.0.0.1:5000 app:app导致只监听 localhost外部无法访问。解决检查Dockerfile中CMD行是否为[gunicorn, --bind, 0.0.0.0:5000, ...]必须是0.0.0.0而非127.0.0.1若已构建镜像用docker commit重新提交修正后的容器。4.5 现象训练train.py时loss一直为nanaccuracy停在 0.25随机猜测水平原因data/目录下train/和val/子目录结构错误。例如train/下直接放图片而非train/可回收物/xxx.jpg或val/里混入了train/的同名图片导致数据泄露。解决执行python -c from torchvision.datasets import ImageFolder; dImageFolder(data/train); print(len(d), d.classes)输出应为(1200, [可回收物, 有害垃圾, 湿垃圾, 干垃圾])若 classes 数量不对用find data/train -type d | grep -v \.$检查子目录名是否含空格或中文标点。5. 模型微调实战如何用你自己的 50 张照片30 分钟追加一个「大骨棒」类别5.1 数据准备比拍照更重要的是文件夹命名规范课程作业默认四分类可回收/有害/湿/干但现实里食堂泔水桶总混进「大骨棒」——它既不是湿垃圾难腐烂也不是干垃圾有油污。想加这个第五类不需要重训整个模型只需增量微调。第一步严格按data/规范新建目录# 在 data/ 下创建新类别目录名称必须是中文与 class_names 一致 mkdir -p data/train/大骨棒 data/val/大骨棒 data/test/大骨棒 # 放入你拍的 50 张大骨棒照片注意不要用手机直接拍的原图 # 先用手机相册批量裁剪确保骨头居中、背景干净 # 再用 Python 脚本批量重命名避免中文乱码 for i in {1..50}; do cp ~/Pictures/大骨棒_$i.jpg data/train/大骨棒/bone_${i}.jpg done提示bone_${i}.jpg这种英文名是安全底线。Windows 资源管理器复制时若含中文Linux 容器内可能读成?????.jpgImageFolder直接跳过该文件。我试过用convmv -f gbk -t utf8 --notest *.jpg转码但不如一开始就用英文名省心。5.2 修改模型与训练脚本两处改动3 分钟搞定model.py中只需改一行扩大输出维度# 原代码4 分类 class ResNet18Classifier(nn.Module): def __init__(self, num_classes4): # ← 改这里 super().__init__() self.resnet models.resnet18(pretrainedTrue) self.resnet.fc nn.Linear(self.resnet.fc.in_features, num_classes) # 改为5 分类 class ResNet18Classifier(nn.Module): def __init__(self, num_classes5): # ← 从 4 改成 5 super().__init__() self.resnet models.resnet18(pretrainedTrue) self.resnet.fc nn.Linear(self.resnet.fc.in_features, num_classes)train.py中改两处一是num_classes二是class_names列表# 原代码 num_classes 4 class_names [可回收物, 有害垃圾, 湿垃圾, 干垃圾] # 改为 num_classes 5 class_names [可回收物, 有害垃圾, 湿垃圾, 干垃圾, 大骨棒] # ← 新增5.3 启动增量训练用旧权重初始化收敛更快关键技巧加载原best_model.pth作为起点只微调最后的全连接层fc layer冻结前面所有层# 修改 train.py 中 model 初始化部分 # 原来是 model ResNet18Classifier(num_classes5) # 改为 model ResNet18Classifier(num_classes5) # 加载旧权重会自动忽略 fc 层不匹配的部分 model.load_state_dict(torch.load(weights/best_model.pth), strictFalse) # 冻结所有层只训练 fc for param in model.parameters(): param.requires_grad False model.resnet.fc.requires_grad True # 只放开 fc 层 # 然后正常 run train.py python train.py这样做的效果是前 49 个 epoch 的 loss 从 2.3 快速降到 0.15第 50 个 epoch 准确率就达 91%。对比从头训练需 200 epochloss 从 1.8 降到 0.4速度提升 4 倍。这是因为 ResNet18 的浅层特征边缘、纹理对所有垃圾都通用只需调整顶层分类器适配新类别。5.4 验证新模型用test/目录做盲测拒绝主观判断别信自己眼睛用代码验证# 新建 test_bone.py import os from predict import predict_image bone_dir data/test/大骨棒 correct 0 total 0 for img_name in os.listdir(bone_dir): if img_name.lower().endswith((.jpg, .jpeg, .png)): result predict_image(os.path.join(bone_dir, img_name)) if result[class] 大骨棒: correct 1 total 1 print(f大骨棒测试集准确率: {correct/total:.3f} ({correct}/{total}))运行python test_bone.py若输出大骨棒测试集准确率: 0.940 (47/50)说明新增类别已稳定生效。这时再把weights/best_model.pth复制到app.py同级目录重启 Flask 或 Docker网页上传大骨棒照片就能正确识别。从那以后我每次加新类别都强制走一遍test_bone.py盲测流程哪怕只是加一张图。因为助教答辩时最爱问“你这个新类别测试集准确率多少有没有混淆到其他类”——没有数字只有“我觉得挺准的”分数立刻掉到 80 分档。希望帮到你。本文还有配套的精品资源点击获取
RELATED READING

延伸阅读

更多一线实战笔记与深度复盘,助您持续精进