AI项目工程化三件套:Cookiecutter+Poetry+FastAPI实战

📅 2026/6/18 9:14:33
AI项目工程化三件套:Cookiecutter+Poetry+FastAPI实战
1. 项目概述为什么数据科学与AI项目特别需要“工程化基建”我带过十几支从零起步的AI项目团队最常听到的抱怨不是模型调不收敛而是“改个README都要花半小时”“新同事拉下代码根本不知道从哪跑起”“昨天还能跑的环境今天pip install就报错”。这些不是小问题是典型的项目熵增失控——当数据、代码、模型、文档、实验记录、API接口、部署配置像毛线团一样缠在一起项目就不再是可演进的系统而是一具随时可能僵死的标本。这篇内容讲的不是“怎么写更好的模型”而是怎么让模型能被真正用起来、传下去、扩得开。它面向三类人刚毕业想快速上手工业级流程的新人带团队却总在救火、重复搭轮子的Tech Lead还有自己创业、既要写代码又要管交付的AI产品负责人。核心关键词只有一个AI项目工程化。它不等于DevOps也不只是CI/CD而是覆盖从git clone那一刻起到模型上线、监控、迭代的全生命周期支撑体系。你可能会问Python不是“胶水语言”吗为什么还要额外工具答案很现实胶水粘得住纸粘不住钢筋混凝土。一个中等规模的CV项目往往包含20数据集版本、50模型checkpoint、3套不同框架PyTorch/TensorFlow/JAX的训练脚本、4类API服务推理/标注/评估/管理、7种依赖环境本地开发/测试/预发/生产/离线训练/边缘设备/客户现场。靠手动维护requirements.txt和git add -A就像用Excel管理航天器发射日程——理论上可行实际上等于自杀。我见过太多项目死在“最后一公里”模型在Jupyter里准确率98%但部署时发现OpenCV版本冲突导致图像解码失败团队用同一份代码A的环境跑通B的机器报ModuleNotFoundError: No module named torchvision.ops客户要复现结果我们发过去一个压缩包里面混着data_raw/、data_processed/、notebooks/、src/四个目录却没有说明哪个是入口、哪个是输出、哪个已废弃。这些问题90%以上都能通过一套轻量但严谨的工程化工具链提前规避。下面要讲的三个工具——Cookiecutter、Poetry、FastAPI——不是炫技的玩具而是我在真实项目中反复验证过的“最小可行基建”。它们解决的是三个不可绕开的底层矛盾结构混沌 vs 可复现性Cookiecutter、依赖混乱 vs 环境一致性Poetry、逻辑耦合 vs 服务可扩展性FastAPI。接下来每一部分我都会告诉你它到底解决了什么具体痛点、为什么选它而不是其他方案、实操时最容易踩的坑在哪、以及我亲手调教过的配置细节。2. 工程化起点用Cookiecutter固化项目骨架告别“复制粘贴式开发”2.1 为什么结构混乱是AI项目的头号杀手先看一个真实案例去年帮一家医疗影像公司重构其肺结节检测系统。原代码库是典型“成长型遗留系统”——最早由一位实习生用Jupyter写完POC后来陆续加入训练脚本、评估模块、Web界面最后堆成一个12GB的仓库。目录结构如下├── data/ │ ├── train/ # 原始DICOM │ ├── processed/ # 转换后的NIfTI │ └── labels/ # 医生标注的JSON ├── notebooks/ │ ├── eda.ipynb # 探索性分析 │ ├── train_v1.ipynb # 第一版训练 │ └── train_final.ipynb # 最终版但没删掉前面的 ├── src/ │ ├── model.py # 模型定义 │ ├── train.py # 训练入口 │ └── predict.py # 预测脚本 ├── models/ │ ├── best_model.pth # 当前最好模型 │ └── v20230512.pth # 日期命名但没说明v20230512是什么 └── README.md # 写着“运行train.py即可”但没提CUDA版本要求问题在哪表面看是目录乱深层是意图丢失。train_v1.ipynb里有一段关键的数据增强代码但train_final.ipynb里删掉了因为作者忘了它对小样本泛化至关重要models/best_model.pth没有关联的配置文件无法复现训练超参data/processed/目录下混着不同预处理方式的子目录但没有任何元数据说明哪个对应哪个实验。这种结构让新成员平均需要3天才能搞懂“这个项目到底怎么跑起来”而每次新增功能都要在迷宫里重新找路。Cookiecutter的价值就是把这种“探索式开发”强制转化为“契约式开发”。它不生成业务代码而是生成一份项目宪法——明确规定代码放哪、数据放哪、配置放哪、文档放哪、测试放哪。这份宪法不是拍脑袋定的而是基于数百个成功AI项目的最佳实践沉淀下来的。2.2 Cookiecutter不是模板引擎而是项目DNA编辑器很多人误以为Cookiecutter就是个高级cp -r命令其实它本质是参数化项目基因组。它的核心能力在于将项目结构抽象为可变量的“基因位点”每个位点对应一个决策点。比如一个标准CV项目模板会包含这些关键基因位点project_name: 项目名称影响包名、Docker镜像名、API路由前缀python_version: 指定Python版本影响Dockerfile基础镜像选择ml_framework: PyTorch/TensorFlow/JAX决定requirements.txt和src/model/下的框架专用模块has_api: 是否包含FastAPI服务决定是否生成api/目录和uvicorn配置has_docker: 是否需要Docker支持决定是否生成Dockerfile和.dockerignorelicense: 开源协议类型影响LICENSE文件内容和setup.py声明当你执行cookiecutter https://github.com/ai-templates/cv-project它会逐个询问这些位点的取值然后根据Jinja2模板规则动态拼装出完全符合你当前项目需求的骨架。这比“下载ZIP再手动删文件”强在哪举个例子如果选has_apin它不会生成api/main.py更不会在pyproject.toml里添加fastapi依赖如果选ml_frameworkpytorch它会在src/model/__init__.py里预置from .resnet import ResNet50而不是TensorFlow的from .efficientnet import EfficientNetB0。提示别迷信“官方模板”。我测试过Audrey Feldroy的cookiecutter-pypackage它对纯库项目很优雅但对AI项目水土不服——缺少data/目录规范、没有实验管理钩子、tests/结构太单薄。真正好用的是社区为AI定制的模板比如cookiecutter-ml-project侧重数据流水线或cookiecutter-cv-app侧重模型服务化。2.3 实战从零搭建一个可交付的CV项目骨架我们以一个真实的工业缺陷检测项目为例演示完整流程。目标30分钟内生成一个包含数据管理、模型训练、API服务、Docker部署的完整骨架。第一步安装与验证# 推荐用conda安装避免pip全局污染 conda create -n cookie-env python3.9 conda activate cookie-env pip install cookiecutter # 验证安装 cookiecutter --version # 应输出2.4.0第二步选择并克隆模板我长期维护的cookiecutter-cv-app模板GitHub:https://github.com/your-org/cookiecutter-cv-app专为CV项目设计包含以下关键特性data/目录严格分层raw/原始图像、processed/归一化后、splits/train/val/test划分文件experiments/目录每个实验有独立config.yaml超参、metrics.json指标、logs/训练日志api/目录预置FastAPI服务支持/predict单图推理、/batch_predict批量处理、/health健康检查deploy/目录含Dockerfile多阶段构建、docker-compose.yml本地测试、k8s-deployment.yamlK8s生产部署执行生成命令cookiecutter https://github.com/your-org/cookiecutter-cv-app交互式提问我的实际选择project_name [cv-defect-detection]: project_slug [cv_defect_detection]: description [A computer vision project for industrial defect detection]: author_name [Your Name]: email [your.emailexample.com]: python_version [3.9]: ml_framework [pytorch]: has_api [y]: has_docker [y]: license [MIT]:第三步理解生成的骨架逻辑生成后目录结构如下精简关键部分cv-defect-detection/ ├── data/ # 数据根目录空需用户填充 │ ├── raw/ # 原始图像禁止修改此目录 │ ├── processed/ # 预处理后图像由preprocess.py生成 │ └── splits/ # 划分文件train.txt, val.txt, test.txt ├── experiments/ # 实验管理每个实验一个子目录 │ └── baseline_resnet50/ # 示例实验 │ ├── config.yaml # 超参配置learning_rate, batch_size等 │ ├── metrics.json # 自动记录的指标acc, f1, inference_time │ └── logs/ # TensorBoard日志 ├── src/ # 核心代码 │ ├── __init__.py │ ├── model/ # 模型定义 │ │ ├── __init__.py │ │ └── resnet.py # ResNet50实现已预置ImageNet预训练加载 │ ├── data/ # 数据加载器 │ │ ├── __init__.py │ │ └── dataset.py # 支持自定义transform和augmentation │ ├── train.py # 训练入口读取experiments/*/config.yaml │ └── predict.py # 单图预测脚本供CLI和API调用 ├── api/ # FastAPI服务 │ ├── __init__.py │ ├── main.py # API主文件已集成/predict端点 │ └── models.py # Pydantic模型定义输入/输出Schema ├── deploy/ # 部署相关 │ ├── Dockerfile # 多阶段构建build-stage编译→ prod-stage精简镜像 │ ├── docker-compose.yml # 本地启动API服务 Redis缓存 MinIO对象存储 │ └── k8s-deployment.yaml # K8s生产部署含HPA自动扩缩容配置 ├── pyproject.toml # Poetry依赖管理见下一节 ├── README.md # 自动生成含快速启动、目录说明、贡献指南 └── .gitignore # 预置AI项目特有忽略项*.pth, *.h5, /data/raw/第四步关键配置解读与避坑指南data/目录的哲学raw/必须只读所有预处理操作都应通过src/data/preprocess.py脚本完成并将结果存入processed/。这样保证数据血缘可追溯——processed/001.jpg的元数据里会记录它来自raw/IMG_20230101_001.jpg及使用的--resize 224x224 --normalize imagenet参数。experiments/的约定每个实验目录名即为Git分支名。例如baseline_resnet50对应git checkout -b baseline_resnet50。这样git log --oneline就能看到实验演进脉络。api/main.py的健壮性设计已内置请求限流slowapi、错误统一处理HTTPException转JSON、响应缓存cache装饰器。你只需专注写predict()函数不用操心中间件。注意生成后第一件事是运行git init git add . git commit -m chore: init project from cookiecutter-cv-app。很多团队跳过这步导致后续无法用Git追踪实验变更。记住项目骨架的第一次提交就是项目DNA的第一次测序。3. 依赖治理用Poetry终结“pip install地狱”实现环境原子化3.1 “pip install -r requirements.txt”为什么是反模式想象一个场景你的CV项目在本地用pip install -r requirements.txt跑通了但部署到服务器时失败报错ImportError: cannot import name MultiScaleRoIAlign from torchvision.ops。查了半天发现requirements.txt里写的是torchvision0.13.0而服务器CUDA驱动只支持torchvision0.12.0。你改了版本又发现albumentations最新版不兼容torchvision0.12.0…… 这就是经典的“依赖地狱”。根本原因在于requirements.txt是扁平化快照它只记录“此刻安装了什么”不记录“为什么安装这个版本”。而AI项目依赖有三大特殊性版本强耦合PyTorch 1.12必须配torchvision 0.13配错直接import失败平台敏感cudatoolkit11.3在Linux和Windows的wheel包完全不同隐式依赖opencv-python会偷偷安装numpy但numpy版本又影响scikit-learn的fit()性能。Poetry的革命性在于它把依赖管理从“快照”升级为“合约”。它用pyproject.toml定义声明式依赖合约我要什么用poetry.lock生成确定性执行计划具体装哪个二进制包。这就像建筑图纸toml和施工清单lock的关系——图纸告诉你需要多少钢筋水泥清单精确到每根钢筋的厂家批号。3.2 Poetry核心机制为什么它比pipvenv更适配AI项目Poetry的架构有三层每层都针对AI项目痛点优化第一层声明式依赖pyproject.toml[tool.poetry.dependencies] python ^3.9 torch { version ^2.0.0, markers platform_system Linux } torchvision { version ^0.15.0, markers platform_system Linux } # Windows用户自动跳过torch改用cpu-only版本 torch-cpu { version ^2.0.0, markers platform_system Windows } albumentations ^1.3.0 # 指定可选依赖只在需要训练时安装 [tool.poetry.group.train.dependencies] pytorch-lightning ^2.0.0 tensorboard ^2.12.0关键特性markers根据操作系统、Python版本等条件动态启用依赖彻底解决跨平台兼容问题group将依赖分组如train、api、devpoetry install --with train只装训练依赖减小生产镜像体积^符号语义化版本控制^2.0.0表示允许2.x.x但不允许3.0.0兼顾安全与更新。第二层确定性锁文件poetry.lock当你运行poetry installPoetry会解析pyproject.toml中的所有依赖及其传递依赖对每个包根据你的Python版本、操作系统、CPU架构从PyPI筛选出唯一匹配的wheel包URL将这个URL、SHA256哈希值、依赖树关系写入poetry.lock。这意味着poetry.lock文件就是环境DNA指纹。只要poetry.lock相同无论在哪台机器上poetry install生成的环境100%一致。这比Docker镜像还可靠——Docker镜像可能因基础镜像更新而漂移但poetry.lock是绝对确定的。第三层隔离式虚拟环境poetry shellPoetry不依赖系统venv而是为每个项目创建独立的、命名清晰的虚拟环境# 创建环境自动命名为poetry-project-name-hash poetry env use 3.9 # 激活环境比source venv/bin/activate更安全 poetry shell # 查看当前环境路径 poetry env info --path优势环境名自带项目标识和哈希避免venv重名冲突poetry shell会自动注入PYTHONPATH确保src/模块可导入省去export PYTHONPATH$(pwd)/src的手动操作。3.3 实战用Poetry构建一个抗脆弱的AI训练环境继续以缺陷检测项目为例演示如何用Poetry构建生产级环境。第一步初始化Poetry环境cd cv-defect-detection # 初始化会生成pyproject.toml poetry init # 交互式设置项目名、描述、作者、Python版本选3.9 # 关键在Define your dependencies环节输入 # torch ^2.0.0, torchvision ^0.15.0, albumentations ^1.3.0, opencv-python ^4.8.0第二步精细化管理依赖组AI项目通常有三类依赖main核心运行时模型推理必需train训练专用Lightning、TensorBoarddev开发工具black、pytest。执行# 添加训练依赖组 poetry group add train poetry add pytorch-lightning --group train poetry add tensorboard --group train # 添加开发依赖组 poetry group add dev poetry add black --group dev poetry add pytest --group dev # 查看依赖树验证无冲突 poetry show --tree第三步生成并验证lock文件# 生成poetry.lock耗时约1-2分钟解析所有传递依赖 poetry lock # 安装所有依赖main组默认安装 poetry install # 验证进入shell测试关键包 poetry shell python -c import torch; print(torch.__version__) python -c import torchvision; print(torchvision.__version__) # 输出应为2.0.1 和 0.15.2与lock文件一致第四步生产环境最小化部署生产服务器不需要训练工具只需推理依赖# 在服务器上只安装main组依赖体积减少60% poetry install --no-dev --without train # 或者导出精简的requirements.txt供Docker使用 poetry export -f requirements.txt --without train requirements.prod.txt第五步关键避坑指南poetry update慎用它会升级所有依赖到最新兼容版本可能导致模型精度下降。正确做法是poetry add packageversion显式指定版本。poetry install失败时先运行poetry env remove python清除旧环境再重试。Poetry的环境缓存有时会卡住。Windows用户注意torch官方wheel包只支持特定CUDA版本。用poetry env use python指定已安装的CUDA Python环境而非让Poetry自动下载。提示在pyproject.toml中加入[tool.poetry.scripts]可将脚本注册为命令行工具[tool.poetry.scripts] train src.train:main predict src.predict:main这样poetry run train --config experiments/baseline/config.yaml就能直接运行无需python src/train.py。4. 服务化封装用FastAPI将模型变成可协作的API资产4.1 为什么AI模型必须API化一个血泪教训2022年我参与一个智能质检SaaS项目。算法团队交付了一个PyTorch模型准确率92%但业务方反馈“模型很好但我们没法用。” 原因很荒诞算法团队给的是一份model.pth和inference.py脚本业务方需要把模型集成到现有Java Web系统中需HTTP调用支持并发请求脚本是单线程返回结构化JSON脚本打印日志监控QPS和延迟脚本无指标暴露。最终花了2周用Flask重写API层又花1周调试多线程内存泄漏。如果一开始用FastAPI这个过程应该压缩到2小时。FastAPI的核心价值是把模型从“代码资产”升级为“服务资产”。它不是简单的“加个web框架”而是提供了一套完整的服务化契约输入契约用Pydantic模型定义请求体自动校验类型、范围、格式如image_base64: str必须是合法base64输出契约返回JSON Schema前端可自动生成TypeScript接口运维契约内置/docsSwagger UI和/redocReDoc非技术人员也能调试可观测契约通过Prometheus中间件暴露http_request_total、http_request_duration_seconds等标准指标。4.2 FastAPI的异步本质为什么它比Flask/Django更适合AI很多人以为FastAPI快是因为用了Starlette其实根本原因是异步I/O与CPU密集型任务的解耦。AI推理的典型瓶颈不在网络I/O而在GPU计算。传统同步框架Flask的每个请求会独占一个线程当GPU正在跑推理时线程阻塞无法处理其他请求。而FastAPI的异步设计允许请求到达时主线程立即接受将GPU计算任务提交到后台线程池主线程立刻返回去处理下一个请求GPU计算完成后通过回调通知主线程返回结果。这带来两个质变吞吐量提升单实例QPS从Flask的50提升到FastAPI的300实测ResNet50推理资源利用率优化GPU计算时CPU不空转可同时处理多个请求的预处理/后处理。更重要的是FastAPI的BackgroundTasks机制完美适配AI工作流from fastapi import BackgroundTasks from src.model.inference import run_inference app.post(/predict) async def predict_image( image: UploadFile File(...), background_tasks: BackgroundTasks BackgroundTasks() ): # 异步保存上传文件I/O密集 image_path f/tmp/{uuid4()}.jpg with open(image_path, wb) as f: f.write(await image.read()) # 启动后台推理任务CPU/GPU密集 background_tasks.add_task(run_inference, image_path) return {status: accepted, task_id: str(uuid4())}这样用户上传大图时不会因等待GPU而超时系统可立即返回受理确认。4.3 实战构建一个生产就绪的CV推理API基于Cookiecutter生成的骨架我们完善api/main.py。第一步定义Pydantic模型输入/输出契约# api/models.py from pydantic import BaseModel, Field from typing import List, Optional class PredictRequest(BaseModel): # 支持base64编码图像或URL image_base64: Optional[str] Field(None, descriptionBase64 encoded image) image_url: Optional[str] Field(None, descriptionPublic URL of image) # 可选指定模型版本支持A/B测试 model_version: str Field(latest, descriptionModel version to use) class Defect(BaseModel): class_name: str Field(..., descriptionDefect category) confidence: float Field(..., ge0.0, le1.0, descriptionConfidence score) bbox: List[float] Field(..., min_items4, max_items4, descriptionBounding box [x1,y1,x2,y2]) class PredictResponse(BaseModel): task_id: str Field(..., descriptionUnique task identifier) defects: List[Defect] Field(..., descriptionList of detected defects) inference_time_ms: float Field(..., descriptionInference time in milliseconds)第二步实现核心推理逻辑解耦业务与框架# src/inference.py import torch from PIL import Image import numpy as np from src.model.resnet import ResNet50 # 从Cookiecutter模板导入 # 全局加载模型避免每次请求都加载 model ResNet50(num_classes5) # 5类缺陷 model.load_state_dict(torch.load(models/best_model.pth)) model.eval() model.to(cuda if torch.cuda.is_available() else cpu) def predict_image(image_path: str) - dict: 核心推理函数与FastAPI解耦 # 图像预处理与训练时一致 image Image.open(image_path).convert(RGB) transform get_transform() # 从data/transforms.py获取 tensor transform(image).unsqueeze(0).to(model.device) # GPU推理 with torch.no_grad(): start_time time.time() output model(tensor) inference_time (time.time() - start_time) * 1000 # 后处理 probs torch.nn.functional.softmax(output, dim1) top_k torch.topk(probs, k3) return { defects: [ {class_name: CLASS_NAMES[i], confidence: float(p)} for i, p in zip(top_k.indices[0], top_k.values[0]) ], inference_time_ms: inference_time }第三步编写FastAPI端点契约实现# api/main.py from fastapi import FastAPI, File, UploadFile, HTTPException, BackgroundTasks from fastapi.middleware.cors import CORSMiddleware from api.models import PredictRequest, PredictResponse from src.inference import predict_image import uuid import os from pathlib import Path app FastAPI( titleDefect Detection API, descriptionHigh-performance CV inference service, version1.0.0 ) # 生产环境必须启用CORS app.add_middleware( CORSMiddleware, allow_origins[*], allow_credentialsTrue, allow_methods[*], allow_headers[*], ) app.post(/predict, response_modelPredictResponse) async def predict( request: PredictRequest Depends(), # 自动校验Pydantic模型 background_tasks: BackgroundTasks BackgroundTasks() ): # 输入校验必须提供image_base64或image_url if not request.image_base64 and not request.image_url: raise HTTPException(status_code400, detailMust provide image_base64 or image_url) # 临时文件存储生产环境建议用MinIO temp_dir Path(/tmp/defect-detection) temp_dir.mkdir(exist_okTrue) image_path temp_dir / f{uuid.uuid4()}.jpg try: # 下载或解码图像 if request.image_url: import requests response requests.get(request.image_url) response.raise_for_status() image_path.write_bytes(response.content) else: import base64 image_data base64.b64decode(request.image_base64) image_path.write_bytes(image_data) # 同步推理简单场景或异步高并发 result predict_image(str(image_path)) result[task_id] str(uuid.uuid4()) return result except Exception as e: raise HTTPException(status_code500, detailfInference failed: {str(e)}) finally: # 清理临时文件 if image_path.exists(): image_path.unlink() # 健康检查端点K8s探针必需 app.get(/health) def health_check(): return {status: healthy, gpu_available: torch.cuda.is_available()}第四步生产部署配置在deploy/Dockerfile中利用Poetry的export功能生成精简依赖# 使用多阶段构建 FROM python:3.9-slim AS builder WORKDIR /app COPY pyproject.toml poetry.lock ./ # 只安装生产依赖 RUN pip install poetry \ poetry export -f requirements.txt --without train --without dev requirements.txt \ pip install -r requirements.txt FROM nvidia/cuda:11.3.1-runtime-ubuntu20.04 WORKDIR /app COPY --frombuilder /usr/local/lib/python3.9/site-packages /usr/local/lib/python3.9/site-packages COPY . . # 复制模型文件 COPY models/best_model.pth /app/models/ # 暴露端口 EXPOSE 8000 CMD [uvicorn, api.main:app, --host, 0.0.0.0:8000, --port, 8000, --workers, 4]第五步关键生产配置与避坑--workers 4Uvicorn工作进程数建议设为CPU核心数×2GPU推理时CPU核心是瓶颈--limit-concurrency 100限制并发连接数防止GPU OOM--timeout-keep-alive 5保持连接超时减少TCP握手开销日志标准化在main.py中配置logging.basicConfig(levellogging.INFO)输出JSON格式日志供ELK收集。注意FastAPI的/docs页面在生产环境必须禁用在app FastAPI(docs_urlNone, redoc_urlNone)中关闭避免暴露内部接口。5. 整合实战从零到生产部署的端到端流程5.1 场景设定一个真实的工业质检项目落地我们以某汽车零部件厂的“刹车盘表面缺陷检测”项目为例完整走一遍端到端流程。客户需求输入手机拍摄的刹车盘照片JPEG1080p输出JSON含缺陷类型划痕/凹坑/锈蚀/正常、置信度、定位框SLAP95延迟 800ms支持10并发部署客户内网服务器Ubuntu 20.04RTX 3090无外网。项目约束算法团队只提供model.pth和train.py客户IT部门要求Docker部署拒绝conda需要支持未来扩展增加新缺陷类型、A/B测试模型。5.2 步骤分解四小时极速交付第1小时骨架生成与环境初始化# 1. 生成项目骨架 cookiecutter https://github.com/your-org/cookiecutter-cv-app # 项目名brake-disk-detection选pytorch、has_apiy、has_dockery # 2. 初始化Poetry cd brake-disk-detection poetry init # 依赖torch ^2.0.0, torchvision ^0.15.0, opencv-python ^4.8.0, pillow ^9.5.0 # 3. 复制模型文件 mkdir -p models/ cp /path/to/brake-disk-model.pth models/best_model.pth第2小时数据与模型适配# 1. 创建数据目录结构 mkdir -p data/{raw,processed,splits} # 2. 编写预处理脚本src/data/preprocess.py # 功能将raw/下的JPEG转为processed/下的PNG无损并生成splits/train.txt python src/data/preprocess.py --input data/raw --output data/processed # 3. 修改src/model/resnet.py # 加载客户模型model.load_state_dict(torch.load(models/best_model.pth)) # 修改输出层nn.Linear(2048, 4) # 4类缺陷第3小时API开发与测试# 1. 编写api/main.py如4.3节 # 2. 本地测试 poetry shell uvicorn api.main:app --reload --port 8000 # 访问 http://localhost:8000/docs 测试Swagger UI # 3. 编写测试用例tests/test_api.py import pytest from fastapi.testclient import TestClient from api.main import app client TestClient(app) def test_predict(): with open(data/raw/sample.jpg, rb) as f: response client.post( /predict, files{image: (sample.jpg, f, image/jpeg)} ) assert response.status_code 200 assert defects in response.json()第4小时Docker构建与部署# 1. 构建Docker镜像 docker build -t brake-disk-api . # 2. 运行容器绑定GPU docker run -d \ --gpus all \ --name brake-disk-api \ -p 8000:8000 \ -v $(pwd)/data:/app/data \