Python AI教程：快速用PyTorch+FastAPI部署手写数字识别接口

conda create -n pyai python=3.10 conda activate pyai

import torch import torch.nn as nn

class Net(nn.Module):

def __init__(self):
    super().__init__()
    self.fc = nn.Sequential(
        nn.Linear(784, 256),
        nn.ReLU(),
        nn.Linear(256, 10)
    )
def forward(self, x):
    return self.fc(x)

from fastapi import FastAPI, File, UploadFile from PIL import Image import torch, torchvision.transforms as T import io

app = FastAPI() model = torch.load('best.pt', map_location='cpu') model.eval()

transform = T.Compose([

T.Grayscale(),
T.Resize((28, 28)),
T.ToTensor(),
T.Normalize((0.1307,), (0.3081,))

])

@app.post('/predict') async def predict(file: UploadFile = File(...)):

img = Image.open(io.BytesIO(await file.read()))
x = transform(img).unsqueeze(0)
with torch.no_grad():
    logits = model(x)
    pred = logits.argmax(dim=1).item()
    prob = torch.softmax(logits, dim=1)[0, pred].item()
return {'digit': pred, 'confidence': round(prob, 4)}