在Windows上部署PyTorch模型的三种主流方法

import torch
import torchvision.models as models

# 1. 加载训练好的模型 (确保处于评估模式)
model = models.resnet18(weights='imagenet1k_v1') # 示例模型
model.eval()

# 2. 创建示例输入 (用于追踪或脚本化)
# 假设输入是 batch_size=1, 3通道, 224x224的图片
example_input = torch.rand(1, 3, 224, 224)

# 3. 跟踪模式 (tracing) - 适合控制流简单的模型
traced_script_module = torch.jit.trace(model, example_input)

# 或者 脚本模式 (scripting) - 适合有复杂控制流(if/for)的模型
# traced_script_module = torch.jit.script(model)

# 4. 保存模型
traced_script_module.save("resnet18_windows.pt")
print("模型已导出为 resnet18_windows.pt")

import torch
import torchvision.transforms as transforms
from pil import image

# 1. 加载 torchscript 模型
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = torch.jit.load("resnet18_windows.pt")
model.to(device)
model.eval()

# 2. 预处理图片
transform = transforms.compose([
    transforms.resize(256),
    transforms.centercrop(224),
    transforms.totensor(),
    transforms.normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])

# 3. 推理
img = image.open("test_image.jpg").convert("rgb")
input_tensor = transform(img).unsqueeze(0).to(device)

with torch.no_grad():
    output = model(input_tensor)
    print("预测结果:", output.argmax(dim=1).item())

pip install fastapi uvicorn[standard] pillow python-multipart
# 如果还没装 torch
pip install torch torchvision

from fastapi import fastapi, file, uploadfile
import torch
import torchvision.transforms as transforms
from pil import image
import io

app = fastapi()

# 全局加载模型 (避免每次请求都加载)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = torch.jit.load("resnet18_windows.pt") # 使用方案一中导出的模型
model.to(device)
model.eval()

transform = transforms.compose([
    transforms.resize(256),
    transforms.centercrop(224),
    transforms.totensor(),
    transforms.normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])

@app.post("/predict/")
async def predict(file: uploadfile = file(...)):
    # 读取图片
    image_data = await file.read()
    image = image.open(io.bytesio(image_data)).convert("rgb")
    
    # 预处理
    input_tensor = transform(image).unsqueeze(0).to(device)
    
    # 推理
    with torch.no_grad():
        output = model(input_tensor)
        prediction = output.argmax(dim=1).item()
    
    return {"filename": file.filename, "class_id": prediction}

# 启动命令: uvicorn main:app --reload --host 0.0.0.0 --port 8000

uvicorn main:app --host 0.0.0.0 --port 8000

pip install pyinstaller pyside6 # 或者 tkinter (内置)

import sys
import torch
from pyside6.qtwidgets import qapplication, qmainwindow, qpushbutton, qlabel, qfiledialog, qvboxlayout, qwidget
from pyside6.qtgui import qpixmap
from pil import image
import torchvision.transforms as transforms
import io

# 加载模型 (全局)
model = torch.jit.load("resnet18_windows.pt")
model.eval()
transform = transforms.compose([...]) # 同上

class mainwindow(qmainwindow):
    def __init__(self):
        super().__init__()
        self.setwindowtitle("pytorch windows 演示")
        
        layout = qvboxlayout()
        self.label = qlabel("请上传图片")
        self.btn = qpushbutton("选择图片并预测")
        self.btn.clicked.connect(self.predict)
        
        layout.addwidget(self.label)
        layout.addwidget(self.btn)
        
        container = qwidget()
        container.setlayout(layout)
        self.setcentralwidget(container)

    def predict(self):
        file_path, _ = qfiledialog.getopenfilename(self, "选择图片", "", "images (*.png *.jpg)")
        if file_path:
            img = image.open(file_path).convert("rgb")
            input_tensor = transform(img).unsqueeze(0)
            with torch.no_grad():
                out = model(input_tensor)
                res = out.argmax(dim=1).item()
            self.label.settext(f"预测类别 id: {res}")

if __name__ == "__main__":
    app = qapplication(sys.argv)
    window = mainwindow()
    window.show()
    sys.exit(app.exec())

pyinstaller --noconfirm --onefile --windowed --add-data "resnet18_windows.pt;." app.py