深度学习实战：自建图形分类器

固定随机种子
      │
      ▼
生成图形图片 ──► train / valid / test
                              │
                              ▼
                    ShapeDataset + DataLoader
                              │
                              ▼
                     TinyShapeCNN 批训练
                              │
              ┌──验证损失──保存最佳 state_dict
              ▼
         一次测试评估
              │
      ┌─────┼─────┐
      ▼          ▼          ▼
  metrics.json  curves.png  confusion.png
      │
      ▼
predict.py 独立加载与单图推理

shape-vision/
├── data/
│   ├── train/
│   ├── valid/
│   └── test/
├── model_output/
├── prepare_data.py
├── dataset.py
├── model.py
├── train.py
├── predict.py
└── requirements.txt

# prepare_data.py
from pathlib import Path

import numpy as np
from PIL import Image, ImageDraw

SEED = 42
SIZE = 32
SPLITS = {"train": 120, "valid": 30, "test": 30}
CLASSES = ("circle", "square", "triangle")
rng = np.random.default_rng(SEED)


def make_image(kind: str) -> Image.Image:
    background = int(rng.integers(5, 35))
    image = Image.new("L", (SIZE, SIZE), color=background)
    draw = ImageDraw.Draw(image)
    object_size = int(rng.integers(13, 23))
    left = int(rng.integers(3, SIZE - object_size - 2))
    top = int(rng.integers(3, SIZE - object_size - 2))
    right, bottom = left + object_size, top + object_size
    fill = int(rng.integers(180, 256))

    if kind == "circle":
        draw.ellipse((left, top, right, bottom), fill=fill)
    elif kind == "square":
        draw.rectangle((left, top, right, bottom), fill=fill)
    elif kind == "triangle":
        draw.polygon(
            [(left + object_size // 2, top), (left, bottom), (right, bottom)],
            fill=fill,
        )
    else:
        raise ValueError(f"unknown shape: {kind}")

    pixels = np.asarray(image, dtype=np.float32)
    pixels += rng.normal(0, 8, pixels.shape)
    return Image.fromarray(np.clip(pixels, 0, 255).astype(np.uint8), mode="L")


root = Path("data")
for split, count in SPLITS.items():
    for class_name in CLASSES:
        folder = root / split / class_name
        folder.mkdir(parents=True, exist_ok=True)
        for index in range(count):
            make_image(class_name).save(folder / f"{class_name}_{index:04d}.png")

print({split: count * len(CLASSES) for split, count in SPLITS.items()})

# dataset.py
from pathlib import Path

import numpy as np
import torch
from PIL import Image
from torch.utils.data import Dataset

CLASSES = ("circle", "square", "triangle")


class ShapeDataset(Dataset):
    def __init__(self, root: str | Path) -> None:
        self.root = Path(root)
        self.samples: list[tuple[Path, int]] = []
        for label, class_name in enumerate(CLASSES):
            paths = sorted((self.root / class_name).glob("*.png"))
            self.samples.extend((path, label) for path in paths)
        if not self.samples:
            raise ValueError(f"no PNG images found in {self.root}")

    def __len__(self) -> int:
        return len(self.samples)

    def __getitem__(self, index: int) -> tuple[torch.Tensor, int]:
        path, label = self.samples[index]
        with Image.open(path) as image:
            pixels = np.asarray(image.convert("L"), dtype=np.float32) / 255.0
        tensor = torch.from_numpy(pixels).unsqueeze(0)
        return tensor, label

# model.py
import torch
from torch import nn


class TinyShapeCNN(nn.Module):
    def __init__(self, classes: int = 3) -> None:
        super().__init__()
        self.features = nn.Sequential(
            nn.Conv2d(1, 8, 3, padding=1),
            nn.BatchNorm2d(8),
            nn.ReLU(),
            nn.MaxPool2d(2),
            nn.Conv2d(8, 16, 3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2),
            nn.Conv2d(16, 24, 3, padding=1),
            nn.ReLU(),
            nn.AdaptiveAvgPool2d((1, 1)),
        )
        self.classifier = nn.Sequential(
            nn.Flatten(),
            nn.Dropout(0.15),
            nn.Linear(24, classes),
        )

    def forward(self, images: torch.Tensor) -> torch.Tensor:
        return self.classifier(self.features(images))

# train.py
import json
from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np
import torch
from torch import nn
from torch.utils.data import DataLoader

from dataset import CLASSES, ShapeDataset
from model import TinyShapeCNN

SEED = 42
torch.manual_seed(SEED)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
output = Path("model_output")
output.mkdir(exist_ok=True)

loaders = {
    "train": DataLoader(ShapeDataset("data/train"), batch_size=32, shuffle=True),
    "valid": DataLoader(ShapeDataset("data/valid"), batch_size=64, shuffle=False),
    "test": DataLoader(ShapeDataset("data/test"), batch_size=64, shuffle=False),
}
model = TinyShapeCNN(len(CLASSES)).to(device)
loss_fn = nn.CrossEntropyLoss()
optimizer = torch.optim.AdamW(model.parameters(), lr=0.003, weight_decay=1e-4)


def run_epoch(loader, training: bool) -> tuple[float, float, list[int], list[int]]:
    model.train(training)
    total_loss, correct = 0.0, 0
    predictions, targets = [], []
    context = torch.enable_grad() if training else torch.inference_mode()
    with context:
        for images, labels in loader:
            images, labels = images.to(device), labels.to(device)
            if training:
                optimizer.zero_grad()
            logits = model(images)
            loss = loss_fn(logits, labels)
            if training:
                loss.backward()
                optimizer.step()
            batch_predictions = logits.argmax(dim=1)
            total_loss += loss.item() * len(labels)
            correct += (batch_predictions == labels).sum().item()
            predictions.extend(batch_predictions.cpu().tolist())
            targets.extend(labels.cpu().tolist())
    size = len(loader.dataset)
    return total_loss / size, correct / size, predictions, targets


history = {"train_loss": [], "valid_loss": [], "valid_accuracy": []}
best_valid_loss = float("inf")
for epoch in range(1, 21):
    train_loss, _, _, _ = run_epoch(loaders["train"], training=True)
    valid_loss, valid_accuracy, _, _ = run_epoch(loaders["valid"], training=False)
    history["train_loss"].append(train_loss)
    history["valid_loss"].append(valid_loss)
    history["valid_accuracy"].append(valid_accuracy)
    if valid_loss < best_valid_loss:
        best_valid_loss = valid_loss
        torch.save(model.state_dict(), output / "shape_cnn.pt")
    print(epoch, round(train_loss, 4), round(valid_loss, 4), round(valid_accuracy, 3))

model.load_state_dict(torch.load(
    output / "shape_cnn.pt", map_location=device, weights_only=True
))
test_loss, test_accuracy, predictions, targets = run_epoch(
    loaders["test"], training=False
)

matrix = np.zeros((len(CLASSES), len(CLASSES)), dtype=int)
for actual, predicted in zip(targets, predictions):
    matrix[actual, predicted] += 1

metadata = {
    "classes": CLASSES,
    "image_size": [32, 32],
    "channels": 1,
    "seed": SEED,
    "torch_version": torch.__version__,
    "test_loss": test_loss,
    "test_accuracy": test_accuracy,
}
(output / "metadata.json").write_text(
    json.dumps(metadata, ensure_ascii=False, indent=2), encoding="utf-8"
)
(output / "history.json").write_text(
    json.dumps(history, indent=2), encoding="utf-8"
)

figure, axis = plt.subplots(figsize=(6, 4))
axis.plot(history["train_loss"], label="train")
axis.plot(history["valid_loss"], label="valid")
axis.set(xlabel="Epoch", ylabel="Loss", title="Shape classifier training")
axis.legend()
figure.tight_layout()
figure.savefig(output / "training_curve.png", dpi=150)
plt.close(figure)

figure, axis = plt.subplots(figsize=(5, 4))
image = axis.imshow(matrix, cmap="Blues")
for row in range(len(CLASSES)):
    for column in range(len(CLASSES)):
        axis.text(column, row, matrix[row, column], ha="center", va="center")
axis.set_xticks(range(len(CLASSES)), CLASSES, rotation=25)
axis.set_yticks(range(len(CLASSES)), CLASSES)
axis.set(xlabel="Predicted", ylabel="Actual", title="Confusion matrix")
figure.colorbar(image, ax=axis)
figure.tight_layout()
figure.savefig(output / "confusion_matrix.png", dpi=150)
plt.close(figure)
print("test:", round(test_loss, 4), round(test_accuracy, 3))
print(matrix)

# predict.py
import json
import sys
from pathlib import Path

import numpy as np
import torch
from PIL import Image

from model import TinyShapeCNN

if len(sys.argv) != 2:
    raise SystemExit("usage: python predict.py IMAGE_PATH")

output = Path("model_output")
metadata = json.loads((output / "metadata.json").read_text(encoding="utf-8"))
classes = metadata["classes"]
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = TinyShapeCNN(len(classes)).to(device)
model.load_state_dict(torch.load(
    output / "shape_cnn.pt", map_location=device, weights_only=True
))
model.eval()

with Image.open(sys.argv[1]) as image:
    image = image.convert("L").resize(tuple(metadata["image_size"]))
    pixels = np.asarray(image, dtype=np.float32) / 255.0
tensor = torch.from_numpy(pixels).unsqueeze(0).unsqueeze(0).to(device)

with torch.inference_mode():
    probabilities = model(tensor).softmax(dim=1)[0].cpu()
prediction = int(probabilities.argmax())
print("class:", classes[prediction])
print("probabilities:", {
    name: round(float(probability), 4)
    for name, probability in zip(classes, probabilities)
})

torch
numpy
Pillow
matplotlib

python3 -m venv .venv
source .venv/bin/activate
python -m pip install -r requirements.txt
python prepare_data.py
python train.py
python predict.py data/test/circle/circle_0000.png
find model_output -maxdepth 1 -type f -print