📚博客主页：knighthood2001
✨公众号：认知up吧 （目前正在带领大家一起提升认知，感兴趣可以来围观一下）
🎃知识星球：【认知up吧|成长|副业】介绍
❤️如遇文章付费，可先看看我公众号中是否发布免费文章❤️
🙏笔者水平有限，欢迎各位大佬指点，相互学习进步！
假设，你有模型，有训练好的模型文件，有模型推理代码，就可以把他放到flask上进行展示。
项目架构
index.html是模板文件app.py是项目运行的入口best_model.pth是训练好的模型参数model.py是神经网络模型，这里采用的是GoogleNet网络。model_reasoning.py是模型推理，通过这里面的代码，我们可以在本地进行猫狗图片的预测。
运行图
点击选择文件
图片下面就显示预测结果了。
项目完整代码与讲解
index.html
<!DOCTYPE html>
<html lang="en">
<head><meta charset="UTF-8"><title>图像分类</title><style>body {font-family: Arial, sans-serif;margin: 20px;}#result {margin-top: 10px;}#preview-image {max-width: 400px;margin-top: 20px;}</style>
</head>
<body><h1>图像分类</h1><form id="upload-form" action="/predict" method="post" enctype="multipart/form-data"><input type="file" name="file" accept="image/*" onchange="previewImage(event)"><input type="submit" value="预测"></form><img id="preview-image" src="" alt=""><br><div id="result"></div><script>document.getElementById('upload-form').addEventListener('submit', async (e) => {e.preventDefault(); // 阻止默认的表单提交行为const formData = new FormData(); // 创建一个新的FormData对象，用于封装表单数据formData.append('file', document.querySelector('input[type=file]').files[0]); // 添加表单数据// 使用fetch API发送POST请求到'/predict'路径，并将formData作为请求体const response = await fetch('/predict', {method: 'POST',body: formData});// 获取响应的JSON数据const result = await response.json();// 将预测结果显示在页面上ID为'result'的元素中document.getElementById('result').innerText = `预测结果: ${result.prediction}`;});function previewImage(event) {const file = event.target.files[0]; // 获取上传的文件对象const reader = new FileReader(); // 创建一个FileReader对象，用于读取文件内容// 清空上一次的预测结果document.getElementById('result').innerText = '';// 当文件读取完成后，将文件内容显示在页面上ID为'preview-image'的元素中reader.onload = function(event) {document.getElementById('preview-image').setAttribute('src', event.target.result);}// 如果用户选择了文件，则开始读取文件内容if (file) {reader.readAsDataURL(file); // 将文件读取为DataURL格式，这样可以直接用作img元素的src属性}}</script>
</body>
</html>前端我练的不多，很多解释已经在代码中讲了。
model.py
这是GoogleNet的网络架构
import torch
from torch import nn
from torchsummary import summary
# 定义一个Inception模块
class Inception(nn.Module):def __init__(self, in_channels, c1, c2, c3, c4): # 这些参数，所在的位置都会发送变化，所有需要这个参数super(Inception, self).__init__()self.ReLU = nn.ReLU()# 路线1，单1×1卷积层self.p1_1 = nn.Conv2d(in_channels=in_channels, out_channels=c1, kernel_size=1)# 路线2，1×1卷积层, 3×3的卷积self.p2_1 = nn.Conv2d(in_channels=in_channels, out_channels=c2[0], kernel_size=1)self.p2_2 = nn.Conv2d(in_channels=c2[0], out_channels=c2[1], kernel_size=3, padding=1)# 路线3，1×1卷积层, 5×5的卷积self.p3_1 = nn.Conv2d(in_channels=in_channels, out_channels=c3[0], kernel_size=1)self.p3_2 = nn.Conv2d(in_channels=c3[0], out_channels=c3[1], kernel_size=5, padding=2)# 路线4，3×3的最大池化, 1×1的卷积self.p4_1 = nn.MaxPool2d(kernel_size=3, padding=1, stride=1)self.p4_2 = nn.Conv2d(in_channels=in_channels, out_channels=c4, kernel_size=1)def forward(self, x):p1 = self.ReLU(self.p1_1(x))p2 = self.ReLU(self.p2_2(self.ReLU(self.p2_1(x))))p3 = self.ReLU(self.p3_2(self.ReLU(self.p3_1(x))))p4 = self.ReLU(self.p4_2(self.p4_1(x)))return torch.cat((p1, p2, p3, p4), dim=1)class GoogLeNet(nn.Module):def __init__(self, Inception, in_channels, out_channels):super(GoogLeNet, self).__init__()self.b1 = nn.Sequential(nn.Conv2d(in_channels=in_channels, out_channels=64, kernel_size=7, stride=2, padding=3),nn.ReLU(),nn.MaxPool2d(kernel_size=3, stride=2, padding=1))self.b2 = nn.Sequential(nn.Conv2d(in_channels=64, out_channels=64, kernel_size=1),nn.ReLU(),nn.Conv2d(in_channels=64, out_channels=192, kernel_size=3, padding=1),nn.ReLU(),nn.MaxPool2d(kernel_size=3, stride=2, padding=1))self.b3 = nn.Sequential(Inception(192, 64, (96, 128), (16, 32), 32),Inception(256, 128, (128, 192), (32, 96), 64),nn.MaxPool2d(kernel_size=3, stride=2, padding=1))self.b4 = nn.Sequential(Inception(480, 192, (96, 208), (16, 48), 64),Inception(512, 160, (112, 224), (24, 64), 64),Inception(512, 128, (128, 256), (24, 64), 64),Inception(512, 112, (128, 288), (32, 64), 64),Inception(528, 256, (160, 320), (32, 128), 128),nn.MaxPool2d(kernel_size=3, stride=2, padding=1))self.b5 = nn.Sequential(Inception(832, 256, (160, 320), (32, 128), 128),Inception(832, 384, (192, 384), (48, 128), 128),nn.AdaptiveAvgPool2d((1, 1)),nn.Flatten(),nn.Linear(1024, out_channels))for m in self.modules():if isinstance(m, nn.Conv2d):nn.init.kaiming_normal_(m.weight, mode="fan_out", nonlinearity='relu')if m.bias is not None:nn.init.constant_(m.bias, 0)elif isinstance(m, nn.Linear):nn.init.normal_(m.weight, 0, 0.01)if m.bias is not None:nn.init.constant_(m.bias, 0)def forward(self, x):x = self.b1(x)x = self.b2(x)x = self.b3(x)x = self.b4(x)x = self.b5(x)return xif __name__ == "__main__":device = torch.device("cuda" if torch.cuda.is_available() else "cpu")model = GoogLeNet(Inception, 1, 10).to(device)print(summary(model, (1, 224, 224)))
model_reasoning.py
import torch
from torchvision import transforms
from model import GoogLeNet, Inception
from PIL import Imagedef test_model(model, test_file):# 设定测试所用到的设备，有GPU用GPU没有GPU用CPUdevice = "cuda" if torch.cuda.is_available() else 'cpu'model = model.to(device)classes = ['猫', '狗']print(classes)image = Image.open(test_file)# normalize = transforms.Normalize([0.162, 0.151, 0.138], [0.058, 0.052, 0.048])# # 定义数据集处理方法变量# test_transform = transforms.Compose([transforms.Resize((224, 224)), transforms.ToTensor(), normalize])# 定义数据集处理方法变量test_transform = transforms.Compose([transforms.Resize((224, 224)), transforms.ToTensor()])image = test_transform(image)# 添加批次维度，变成[1,3,224,224]image = image.unsqueeze(0)with torch.no_grad():model.eval()image = image.to(device) # 图片也要放到设备当中output = model(image)print(output.tolist())pre_lab = torch.argmax(output, dim=1)result = pre_lab.item()print("预测值：", classes[result])return classes[result]def test_special_model(best_model_file, test_file):# 加载模型model = GoogLeNet(Inception, in_channels=3, out_channels=2)model.load_state_dict(torch.load(best_model_file))# 模型的推理判断return test_model(model, test_file)if __name__ == "__main__":# # 加载模型# model = GoogLeNet(Inception, in_channels=3, out_channels=2)# model.load_state_dict(torch.load('best_model.pth'))# # 模型的推理判断# test_model(model, "test_data/images.jfif")test_special_model("best_model.pth", "static/1.jpg")
这段代码与之前的模型推理代码不同的是，我添加了test_special_model函数，方便后续app.py中可以直接调用这个函数进行模型推理。
app.py
import os
from flask import Flask, request, jsonify, render_templatefrom model_reasoning import test_special_model
from model_reasoning import test_model
app = Flask(__name__)# 定义路由
@app.route('/')
def index():return render_template('index.html')@app.route('/predict', methods=['POST'])
def predict():if request.method == 'POST':# 获取上传的文件file = request.files['file']if file:# 调用模型进行预测# # 加载模型# model = GoogLeNet(Inception, in_channels=3, out_channels=2)# basedir = os.path.abspath(os.path.dirname(__file__))## model.load_state_dict(torch.load(basedir + '/best_model.pth'))# result = test_model(model, file)basedir = os.path.abspath(os.path.dirname(__file__))best_model_file = basedir + '/best_model.pth'result = test_special_model(best_model_file, file)return jsonify({'prediction': result})else:return jsonify({'error': 'No file found'})if __name__ == '__main__':app.run(debug=True)
如果没有上文中的test_special_model函数，那么这里你就需要
# 加载模型model = GoogLeNet(Inception, in_channels=3, out_channels=2)basedir = os.path.abspath(os.path.dirname(__file__))model.load_state_dict(torch.load(basedir + '/best_model.pth'))result = test_model(model, file)
并且还需要导入相应的库。
best_model.pth
最重要的是，你需要训练好的一个模型。
有需要的，可以联系我，我直接把这个项目代码发你。省得你还需要配置项目架构。
小插曲
我为什么会使用绝对路径，因为我在使用相对路径后，代码提示找不到这个路径。
basedir = os.path.abspath(os.path.dirname(__file__))best_model_file = basedir + '/best_model.pth'
然后，我刚刚又试了一下，发现使用相对路径，又可以运行成功了。
真是不可思议（这个小插曲花了我大半个小时）。
