项目实战

# 项目结构示例
image_classification/
├── data/
│   ├── train/
│   ├── val/
│   └── test/
├── models/
│   ├── __init__.py
│   ├── cnn_model.py
│   └── resnet_model.py
├── utils/
│   ├── __init__.py
│   ├── data_loader.py
│   └── transforms.py
├── train.py
├── evaluate.py
├── predict.py
├── app.py
└── requirements.txt

# 主训练脚本
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from models.cnn_model import CNNModel
from utils.data_loader import ImageDataset
from utils.transforms import get_transforms

def train_model(model, train_loader, val_loader, epochs=100):
    """模型训练函数"""
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.Adam(model.parameters(), lr=0.001)
    scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=30, gamma=0.1)
    
    for epoch in range(epochs):
        # 训练阶段
        model.train()
        train_loss = 0.0
        for batch_idx, (data, target) in enumerate(train_loader):
            optimizer.zero_grad()
            output = model(data)
            loss = criterion(output, target)
            loss.backward()
            optimizer.step()
            train_loss += loss.item()
        
        # 验证阶段
        model.eval()
        val_loss = 0.0
        correct = 0
        total = 0
        with torch.no_grad():
            for data, target in val_loader:
                output = model(data)
                val_loss += criterion(output, target).item()
                pred = output.argmax(dim=1, keepdim=True)
                correct += pred.eq(target.view_as(pred)).sum().item()
                total += target.size(0)
        
        scheduler.step()
        
        print(f'Epoch {epoch+1}/{epochs}:')
        print(f'Train Loss: {train_loss/len(train_loader):.4f}')
        print(f'Val Loss: {val_loss/len(val_loader):.4f}')
        print(f'Val Accuracy: {100.*correct/total:.2f}%')

# Flask Web应用
from flask import Flask, request, jsonify
import torch
from PIL import Image
import io

app = Flask(__name__)
model = CNNModel()
model.load_state_dict(torch.load('best_model.pth'))
model.eval()

@app.route('/predict', methods=['POST'])
def predict():
    """预测接口"""
    if 'image' not in request.files:
        return jsonify({'error': 'No image provided'}), 400
    
    file = request.files['image']
    image = Image.open(io.BytesIO(file.read()))
    
    # 预处理图像
    transform = get_transforms()
    image_tensor = transform(image).unsqueeze(0)
    
    # 预测
    with torch.no_grad():
        output = model(image_tensor)
        probabilities = torch.softmax(output, dim=1)
        predicted_class = torch.argmax(probabilities, dim=1).item()
        confidence = probabilities[0][predicted_class].item()
    
    return jsonify({
        'class': predicted_class,
        'confidence': confidence,
        'probabilities': probabilities[0].tolist()
    })

if __name__ == '__main__':
    app.run(debug=True, host='0.0.0.0', port=5000)