在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