一文看懂AI技术栈4层架构5个选型原则一、问题AI 技术栈碎片化如何高效构建当你开始一个 AI 项目迎面而来的是一堆陌生名词Python、CUDA、PyTorch、TensorFlow、ONNX、TensorRT、Docker、Kubernetes……每个环节都有多种可选方案稍有不慎就会陷入“框架切换”“环境冲突”“推理速度慢”等坑里。我见过太多开发者在“学新框架”的路上越走越远却连一个端到端的 MVP 都跑不通。根本矛盾AI 技术栈的广度与深度要求远超传统软件开发。你需要同时掌握数据处理、模型训练、模型导出、推理部署、性能优化等多个域且每个域的工具都在快速迭代。本文将用一个完整的端到端案例PyTorch 训练 → ONNX 导出 → ONNX Runtime 推理串起 AI 技术栈的核心环节并给出可落地的选型建议。—## 二、方案分层拆解 AI 技术栈一个典型的 AI 应用技术栈分为四层| 层级 | 核心任务 | 主流工具 | 选型原则 ||------|----------|----------|----------|| 数据处理 | 清洗、增强、加载 | Pandas, PyTorch DataLoader, TF Dataset, NVIDIA DALI | 与框架兼容IO 密集时用 DALI || 模型训练 | 定义网络、优化、调参 | PyTorch, TensorFlow, JAX | 新手推荐 PyTorch生产环境考虑 TorchServe || 模型导出 | 中间表示、量化 | ONNX, TorchScript, TensorRT, OpenVINO | 跨平台必用 ONNXGPU 加速用 TRT || 推理部署 | 服务化、边缘端 | ONNX Runtime, Triton Inference Server, TFLite, CoreML | 高性能用 Triton轻量用 ONNX Runtime |核心思路训练阶段用 PyTorch灵活导出为 ONNX中立推理阶段用 ONNX Runtime高效。这样既能享受 PyTorch 的开发体验又能获得接近原生的推理性能。—## 三、代码实战PyTorch → ONNX → ONNX Runtime 端到端### 1. 环境准备保证可复现bash# 创建虚拟环境安装必要依赖conda create -n ai_stack python3.9conda activate ai_stackpip install torch2.1.0 torchvision0.16.0 onnx1.15.0 onnxruntime1.17.0 numpy1.26.0### 2. 训练一个简单的分类模型我们使用经典的 MNIST 数据集的子集训练一个两层卷积网络。注意这里故意简化数据量只取 1000 张图方便快速运行。python# train_model.pyimport torchimport torch.nn as nnimport torch.optim as optimfrom torchvision import datasets, transformsfrom torch.utils.data import DataLoader, Subset# 0. 定义网络兼容 ONNX 导出避免使用动态控制流class SimpleCNN(nn.Module): def __init__(self, num_classes10): super().__init__() self.features nn.Sequential( nn.Conv2d(1, 16, kernel_size3, padding1), nn.ReLU(), nn.MaxPool2d(2), nn.Conv2d(16, 32, kernel_size3, padding1), nn.ReLU(), nn.MaxPool2d(2), ) self.classifier nn.Sequential( nn.Flatten(), nn.Linear(32 * 7 * 7, 64), nn.ReLU(), nn.Linear(64, num_classes), ) def forward(self, x): x self.features(x) x self.classifier(x) return x# 1. 加载数据仅取前1000张训练200张测试transform transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))])full_train datasets.MNIST(./data, trainTrue, downloadTrue, transformtransform)full_test datasets.MNIST(./data, trainFalse, downloadTrue, transformtransform)train_loader DataLoader(Subset(full_train, indicesrange(1000)), batch_size32, shuffleTrue)test_loader DataLoader(Subset(full_test, indicesrange(200)), batch_size32, shuffleFalse)# 2. 训练device torch.device(cuda if torch.cuda.is_available() else cpu)model SimpleCNN().to(device)criterion nn.CrossEntropyLoss()optimizer optim.Adam(model.parameters(), lr0.001)print(开始训练...)for epoch in range(3): model.train() 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() # 简单验证 model.eval() correct 0 total 0 with torch.no_grad(): for images, labels in test_loader: images, labels images.to(device), labels.to(device) outputs model(images) _, predicted torch.max(outputs, 1) total labels.size(0) correct (predicted labels).sum().item() print(fEpoch {epoch1}, Test Acc: {100 * correct / total:.2f}%)# 保存 PyTorch 权重torch.save(model.state_dict(), model.pth)print(模型权重保存至 model.pth)### 3. 导出为 ONNXONNX 导出时需要指定输入张量的形状动态批处理使用dynamic_axes。注意导出前必须将模型设置为 eval 模式。python# export_onnx.pyimport torchfrom train_model import SimpleCNN, devicemodel SimpleCNN()model.load_state_dict(torch.load(model.pth, map_locationcpu))model.eval()# 构造 dummy 输入batch_size1, channel1, height28, width28dummy_input torch.randn(1, 1, 28, 28)# 定义动态轴batch 维度可变dynamic_axes { input: {0: batch_size}, output: {0: batch_size}}torch.onnx.export( model, dummy_input, model.onnx, input_names[input], output_names[output], dynamic_axesdynamic_axes, opset_version17, do_constant_foldingTrue)print(ONNX 模型导出成功model.onnx)### 4. 使用 ONNX Runtime 进行推理ONNX Runtime 支持 CPU/GPU且无需安装 PyTorch。下面代码展示加载 ONNX 模型并对单张图片进行推理。python# inference_onnx.pyimport onnxruntime as ortimport numpy as npfrom torchvision import datasets, transformsfrom torch.utils.data import DataLoader, Subsetimport time# 加载测试数据中的一张图仅用于演示transform transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))])full_test datasets.MNIST(./data, trainFalse, downloadTrue, transformtransform)test_loader DataLoader(Subset(full_test, indicesrange(1)), batch_size1)image, label next(iter(test_loader))# 转成 numpy并确保维度顺序为 NCHWinput_np image.numpy().astype(np.float32)# 创建 ONNX Runtime 会话session ort.InferenceSession(model.onnx)input_name session.get_inputs()[0].nameoutput_name session.get_outputs()[0].name# 推理预热_ session.run([output_name], {input_name: input_np})# 正式计时模拟批处理batch_np np.repeat(input_np, 10, axis0) # batch_size10start time.time()for _ in range(100): outputs session.run([output_name], {input_name: batch_np})end time.time()print(fONNX Runtime 推理 100 批 (batch10) 耗时: {(end-start)*1000:.2f} ms)pred np.argmax(outputs[0], axis1)print(f预测标签: {pred[:5]}, 真实标签: {label.item()})### 5. 对比 PyTorch 直接推理可选作为对照我们也可以将导出的 ONNX 模型与 PyTorch 推理速度对比。但注意ONNX Runtime 通常比 PyTorch eager 模式快 1.5~3 倍尤其是 CPU 环境。python# compare.py附加代码不单独运行import torchfrom train_model import SimpleCNN, device# PyTorch 推理model_pt SimpleCNN()model_pt.load_state_dict(torch.load(model.pth, map_locationcpu))model_pt.eval()with torch.no_grad(): start time.time() for _ in range(100): _ model_pt(torch.from_numpy(batch_np)) end time.time()print(fPyTorch 推理 100 批耗时: {(end-start)*1000:.2f} ms)说明上述代码直接粘贴即可运行需先运行train_model.py和export_onnx.py。如果 GPU 可用可安装onnxruntime-gpu获得加速。—## 四、总结与最佳实践### 1. 技术栈选型建议-个人/小团队PyTorch ONNX Runtime 组合最灵活学习成本低。-企业生产训练用 PyTorch部署用 NVIDIA Triton Inference Server支持多框架多模型并行。-边缘端考虑 TensorRTNVIDIA 设备或 OpenVINOIntel 设备或者直接使用 TFLite移动端/嵌入式。-量化加速ONNX Runtime 内置动态、静态量化工具FP16/INT8 推理可提升 2~4 倍速度。### 2. 导出 ONNX 的常见陷阱-动态控制流如if语句依赖输入会导致导出失败需提前改写成静态形式如使用torch.where。-自定义算子ONNX 不支持 PyTorch 所有算子需要注册或 fallback。建议先在torch.onnx.export中设置verifyTrue检查。-输入输出动态形状务必设置dynamic_axes否则导出的模型只能固定 batch 大小。### 3. 可操作建议1.从简单模型开始不要一上来就尝试 YOLO / GPT先跑通上面的 MNIST 案例理解整个链路。2.使用版本锁定ONNX 生态与框架版本强相关建议在requirements.txt中固定 torch、onnx、onnxruntime 版本。3.压测找瓶颈用 ONNX Runtime 的session.get_providers()检查是否真的使用了 GPUCUDAExecutionProvider。CPU 环境下可以尝试IExecutionProvider的enable_cpu_mem_arenaFalse以降低内存碎片。4.拥抱容器化用 Docker 打包运行时环境避免不同机器 CUDA/系统库版本冲突。推荐用nvidia/cuda:11.8-runtime-ubuntu22.04作为基础镜像。### 4. 未来的扩展当你掌握了基础链路可以进一步学习-模型服务化使用FastAPI 包装 ONNX Runtime 推理配合 Docker 对外提供 REST API。-自动化量化ONNX Runtime 的 QAT量化感知训练 在保持精度的同时大幅提速。-多 GPU 推理用 Triton Inference Server 管理多进程推理吞吐量翻倍。AI 技术栈的深度远不止上述内容但“训练→导出→部署”这一闭环是所有 AI 工程师必须掌握的硬技能。不要被碎片化的工具吓倒用最小的闭环启动在实践中扩展认知——这就是最优的成长路径。
AI的技术栈全知道
一文看懂AI技术栈4层架构5个选型原则一、问题AI 技术栈碎片化如何高效构建当你开始一个 AI 项目迎面而来的是一堆陌生名词Python、CUDA、PyTorch、TensorFlow、ONNX、TensorRT、Docker、Kubernetes……每个环节都有多种可选方案稍有不慎就会陷入“框架切换”“环境冲突”“推理速度慢”等坑里。我见过太多开发者在“学新框架”的路上越走越远却连一个端到端的 MVP 都跑不通。根本矛盾AI 技术栈的广度与深度要求远超传统软件开发。你需要同时掌握数据处理、模型训练、模型导出、推理部署、性能优化等多个域且每个域的工具都在快速迭代。本文将用一个完整的端到端案例PyTorch 训练 → ONNX 导出 → ONNX Runtime 推理串起 AI 技术栈的核心环节并给出可落地的选型建议。—## 二、方案分层拆解 AI 技术栈一个典型的 AI 应用技术栈分为四层| 层级 | 核心任务 | 主流工具 | 选型原则 ||------|----------|----------|----------|| 数据处理 | 清洗、增强、加载 | Pandas, PyTorch DataLoader, TF Dataset, NVIDIA DALI | 与框架兼容IO 密集时用 DALI || 模型训练 | 定义网络、优化、调参 | PyTorch, TensorFlow, JAX | 新手推荐 PyTorch生产环境考虑 TorchServe || 模型导出 | 中间表示、量化 | ONNX, TorchScript, TensorRT, OpenVINO | 跨平台必用 ONNXGPU 加速用 TRT || 推理部署 | 服务化、边缘端 | ONNX Runtime, Triton Inference Server, TFLite, CoreML | 高性能用 Triton轻量用 ONNX Runtime |核心思路训练阶段用 PyTorch灵活导出为 ONNX中立推理阶段用 ONNX Runtime高效。这样既能享受 PyTorch 的开发体验又能获得接近原生的推理性能。—## 三、代码实战PyTorch → ONNX → ONNX Runtime 端到端### 1. 环境准备保证可复现bash# 创建虚拟环境安装必要依赖conda create -n ai_stack python3.9conda activate ai_stackpip install torch2.1.0 torchvision0.16.0 onnx1.15.0 onnxruntime1.17.0 numpy1.26.0### 2. 训练一个简单的分类模型我们使用经典的 MNIST 数据集的子集训练一个两层卷积网络。注意这里故意简化数据量只取 1000 张图方便快速运行。python# train_model.pyimport torchimport torch.nn as nnimport torch.optim as optimfrom torchvision import datasets, transformsfrom torch.utils.data import DataLoader, Subset# 0. 定义网络兼容 ONNX 导出避免使用动态控制流class SimpleCNN(nn.Module): def __init__(self, num_classes10): super().__init__() self.features nn.Sequential( nn.Conv2d(1, 16, kernel_size3, padding1), nn.ReLU(), nn.MaxPool2d(2), nn.Conv2d(16, 32, kernel_size3, padding1), nn.ReLU(), nn.MaxPool2d(2), ) self.classifier nn.Sequential( nn.Flatten(), nn.Linear(32 * 7 * 7, 64), nn.ReLU(), nn.Linear(64, num_classes), ) def forward(self, x): x self.features(x) x self.classifier(x) return x# 1. 加载数据仅取前1000张训练200张测试transform transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))])full_train datasets.MNIST(./data, trainTrue, downloadTrue, transformtransform)full_test datasets.MNIST(./data, trainFalse, downloadTrue, transformtransform)train_loader DataLoader(Subset(full_train, indicesrange(1000)), batch_size32, shuffleTrue)test_loader DataLoader(Subset(full_test, indicesrange(200)), batch_size32, shuffleFalse)# 2. 训练device torch.device(cuda if torch.cuda.is_available() else cpu)model SimpleCNN().to(device)criterion nn.CrossEntropyLoss()optimizer optim.Adam(model.parameters(), lr0.001)print(开始训练...)for epoch in range(3): model.train() 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() # 简单验证 model.eval() correct 0 total 0 with torch.no_grad(): for images, labels in test_loader: images, labels images.to(device), labels.to(device) outputs model(images) _, predicted torch.max(outputs, 1) total labels.size(0) correct (predicted labels).sum().item() print(fEpoch {epoch1}, Test Acc: {100 * correct / total:.2f}%)# 保存 PyTorch 权重torch.save(model.state_dict(), model.pth)print(模型权重保存至 model.pth)### 3. 导出为 ONNXONNX 导出时需要指定输入张量的形状动态批处理使用dynamic_axes。注意导出前必须将模型设置为 eval 模式。python# export_onnx.pyimport torchfrom train_model import SimpleCNN, devicemodel SimpleCNN()model.load_state_dict(torch.load(model.pth, map_locationcpu))model.eval()# 构造 dummy 输入batch_size1, channel1, height28, width28dummy_input torch.randn(1, 1, 28, 28)# 定义动态轴batch 维度可变dynamic_axes { input: {0: batch_size}, output: {0: batch_size}}torch.onnx.export( model, dummy_input, model.onnx, input_names[input], output_names[output], dynamic_axesdynamic_axes, opset_version17, do_constant_foldingTrue)print(ONNX 模型导出成功model.onnx)### 4. 使用 ONNX Runtime 进行推理ONNX Runtime 支持 CPU/GPU且无需安装 PyTorch。下面代码展示加载 ONNX 模型并对单张图片进行推理。python# inference_onnx.pyimport onnxruntime as ortimport numpy as npfrom torchvision import datasets, transformsfrom torch.utils.data import DataLoader, Subsetimport time# 加载测试数据中的一张图仅用于演示transform transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))])full_test datasets.MNIST(./data, trainFalse, downloadTrue, transformtransform)test_loader DataLoader(Subset(full_test, indicesrange(1)), batch_size1)image, label next(iter(test_loader))# 转成 numpy并确保维度顺序为 NCHWinput_np image.numpy().astype(np.float32)# 创建 ONNX Runtime 会话session ort.InferenceSession(model.onnx)input_name session.get_inputs()[0].nameoutput_name session.get_outputs()[0].name# 推理预热_ session.run([output_name], {input_name: input_np})# 正式计时模拟批处理batch_np np.repeat(input_np, 10, axis0) # batch_size10start time.time()for _ in range(100): outputs session.run([output_name], {input_name: batch_np})end time.time()print(fONNX Runtime 推理 100 批 (batch10) 耗时: {(end-start)*1000:.2f} ms)pred np.argmax(outputs[0], axis1)print(f预测标签: {pred[:5]}, 真实标签: {label.item()})### 5. 对比 PyTorch 直接推理可选作为对照我们也可以将导出的 ONNX 模型与 PyTorch 推理速度对比。但注意ONNX Runtime 通常比 PyTorch eager 模式快 1.5~3 倍尤其是 CPU 环境。python# compare.py附加代码不单独运行import torchfrom train_model import SimpleCNN, device# PyTorch 推理model_pt SimpleCNN()model_pt.load_state_dict(torch.load(model.pth, map_locationcpu))model_pt.eval()with torch.no_grad(): start time.time() for _ in range(100): _ model_pt(torch.from_numpy(batch_np)) end time.time()print(fPyTorch 推理 100 批耗时: {(end-start)*1000:.2f} ms)说明上述代码直接粘贴即可运行需先运行train_model.py和export_onnx.py。如果 GPU 可用可安装onnxruntime-gpu获得加速。—## 四、总结与最佳实践### 1. 技术栈选型建议-个人/小团队PyTorch ONNX Runtime 组合最灵活学习成本低。-企业生产训练用 PyTorch部署用 NVIDIA Triton Inference Server支持多框架多模型并行。-边缘端考虑 TensorRTNVIDIA 设备或 OpenVINOIntel 设备或者直接使用 TFLite移动端/嵌入式。-量化加速ONNX Runtime 内置动态、静态量化工具FP16/INT8 推理可提升 2~4 倍速度。### 2. 导出 ONNX 的常见陷阱-动态控制流如if语句依赖输入会导致导出失败需提前改写成静态形式如使用torch.where。-自定义算子ONNX 不支持 PyTorch 所有算子需要注册或 fallback。建议先在torch.onnx.export中设置verifyTrue检查。-输入输出动态形状务必设置dynamic_axes否则导出的模型只能固定 batch 大小。### 3. 可操作建议1.从简单模型开始不要一上来就尝试 YOLO / GPT先跑通上面的 MNIST 案例理解整个链路。2.使用版本锁定ONNX 生态与框架版本强相关建议在requirements.txt中固定 torch、onnx、onnxruntime 版本。3.压测找瓶颈用 ONNX Runtime 的session.get_providers()检查是否真的使用了 GPUCUDAExecutionProvider。CPU 环境下可以尝试IExecutionProvider的enable_cpu_mem_arenaFalse以降低内存碎片。4.拥抱容器化用 Docker 打包运行时环境避免不同机器 CUDA/系统库版本冲突。推荐用nvidia/cuda:11.8-runtime-ubuntu22.04作为基础镜像。### 4. 未来的扩展当你掌握了基础链路可以进一步学习-模型服务化使用FastAPI 包装 ONNX Runtime 推理配合 Docker 对外提供 REST API。-自动化量化ONNX Runtime 的 QAT量化感知训练 在保持精度的同时大幅提速。-多 GPU 推理用 Triton Inference Server 管理多进程推理吞吐量翻倍。AI 技术栈的深度远不止上述内容但“训练→导出→部署”这一闭环是所有 AI 工程师必须掌握的硬技能。不要被碎片化的工具吓倒用最小的闭环启动在实践中扩展认知——这就是最优的成长路径。
)
双梁正问题 【 torch 实战】（Python代码实现）)
