基于GpuGeek平台的深度学习项目

import os
import cv2
import numpy as np
from sklearn.model_selection import train_test_split
from torchvision import transforms

IMG_SIZE = 48
root_dir = '/dataset/GTSRB'
data, labels = [], []

for class_id in os.listdir(root_dir):
    class_path = os.path.join(root_dir, class_id)
    for img_file in os.listdir(class_path):
        img = cv2.imread(os.path.join(class_path, img_file))
        img = cv2.resize(img, (IMG_SIZE, IMG_SIZE))
        data.append(img)
        labels.append(int(class_id))

X = np.array(data)
y = np.array(labels)
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2)

1.2 数据增强

transform_train = transforms.Compose([
    transforms.ToPILImage(),
    transforms.RandomRotation(15),
    transforms.RandomHorizontalFlip(),
    transforms.ToTensor()
])

2. 模型构建与训练

2.1 使用预训练ResNet50

import torch
import torch.nn as nn
import torchvision.models as models

model = models.resnet50(pretrained=True)
model.fc = nn.Linear(model.fc.in_features, 43)  # 43类交通标志
model = model.cuda()

2.2 模型训练流程

import torch.optim as optim
from torch.utils.data import DataLoader, TensorDataset

criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.0001)

dataloader = DataLoader(TensorDataset(torch.tensor(X_train).permute(0, 3, 1, 2).float() / 255, torch.tensor(y_train)), batch_size=64, shuffle=True)

for epoch in range(10):
    model.train()
    total_loss = 0
    for batch_x, batch_y in dataloader:
        batch_x, batch_y = batch_x.cuda(), batch_y.cuda()
        optimizer.zero_grad()
        outputs = model(batch_x)
        loss = criterion(outputs, batch_y)
        loss.backward()
        optimizer.step()
        total_loss += loss.item()
    print(f"Epoch {epoch+1}, Loss: {total_loss:.4f}")

3. 模型评估

from sklearn.metrics import classification_report

testloader = DataLoader(TensorDataset(torch.tensor(X_test).permute(0, 3, 1, 2).float() / 255, torch.tensor(y_test)), batch_size=64)
model.eval()
y_pred, y_true = [], []

with torch.no_grad():
    for bx, by in testloader:
        bx = bx.cuda()
        outputs = model(bx)
        preds = outputs.argmax(dim=1).cpu().numpy()
        y_pred.extend(preds)
        y_true.extend(by.numpy())

print(classification_report(y_true, y_pred))

四、部署与上线测试

1. 保存与导出模型

torch.save(model.state_dict(), 'traffic_sign_model.pt')

2. 使用GpuGeek创建推理实例

登录GpuGeek后台，选择PyTorch镜像，部署模型推理API。
上传模型权重至实例，编写轻量级Flask服务：

from flask import Flask, request
import torch
from torchvision import transforms

app = Flask(__name__)
model.load_state_dict(torch.load('traffic_sign_model.pt'))
model.eval()

@app.route('/predict', methods=['POST'])
def predict():
    img = request.files['image'].read()
    # 预处理略
    output = model(image_tensor)
    return str(torch.argmax(output).item())

app.run(host='0.0.0.0', port=8000)