用Pytorch实现猫狗二分类

用Pytorch实现猫狗二分类

预训练模型：resnet18 数据集：Kaggle 训练过程：

• 准备数据：下载好Kaggle数据集并按照猫狗分类，可以使用脚本分类
• 数据增强：变换出更多数据，防止过拟合
• 加载预训练模型：加载 resnet18 的预训练权重
• 改造模型：冻结所有层->换掉分类头（1000->2）
• 训练
• 验证
• 保存
• 制作用户界面并使用图片推理

下载的样本按照猫狗分类到两个文件夹里

把图片放到./data/all后运行脚本

import os
import shutil
import random

# ===== Config =====
SRC_DIR = "./data/all"       # 原始图片：cat.*.jpg, dog.*.jpg
DST_DIR = "./data"           # 输出：data/train, data/val
SPLIT_RATIO = 0.8            # 80% 训练，20% 验证
SEED = 42

# ===== Reproducibility =====
random.seed(SEED)

# ===== Split =====
for cls in ["cat", "dog"]:
    # 找所有该类的图
    files = [f for f in os.listdir(SRC_DIR) if f.startswith(cls + ".")]
    random.shuffle(files)

    split = int(len(files) * SPLIT_RATIO)
    train_files = files[:split]
    val_files = files[split:]

    # 建文件夹
    os.makedirs(f"{DST_DIR}/train/{cls}", exist_ok=True)
    os.makedirs(f"{DST_DIR}/val/{cls}", exist_ok=True)

    # 复制文件
    for f in train_files:
        shutil.copy(os.path.join(SRC_DIR, f),
                    os.path.join(DST_DIR, "train", cls, f))
    for f in val_files:
        shutil.copy(os.path.join(SRC_DIR, f),
                    os.path.join(DST_DIR, "val", cls, f))

    print(f"{cls}: train={len(train_files)}, val={len(val_files)}")

print("done!")

导入必要的库＆设置常量

其中MEAN和STD这两个是大量图片色彩的均值与标准差，用于将图片色彩的值转换为较小的数，防止训练中梯度爆炸。

import os
os.environ['TORCH_HOME']='./cache'
import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets,transforms,models

# constants
IMAGENET_MEAN=[0.485,0.456,0.406]
IMAGENET_STD=[0.229,0.224,0.225]

创建预处理流水线

在训练数据的预处理流水线中

• 重新设置尺寸，缩小尺寸
• 随机裁剪图像的一部分，使图片的尺寸符合模型的输入端，并使每一个batch的输入图像不同
• 以百分之五十的概率水平翻转图像
• 亮度，对比度，饱和度的随机化（0.2表示±0.2调整）
• 将图片转化为tensor
• 用平均值和标准差转换原图片，用于输入模型

在测试集的预处理流水线中，在以上的操作步骤中，去掉了颜色抖动和翻转两项，裁剪方式更改，其余处理相同

# Train transform: with augmentation 
train_transform = transforms.Compose([
    transforms.Resize(256),                       # resize short side to 256
    transforms.RandomCrop(224),                   # random crop to 224
    transforms.RandomHorizontalFlip(p=0.5),       # random horizontal flip
    transforms.ColorJitter(brightness=0.2,
                           contrast=0.2,
                           saturation=0.2),        # random color jitter
    transforms.ToTensor(),                        # to tensor, [0,1]
    transforms.Normalize(IMAGENET_MEAN, IMAGENET_STD),
])

# Test transform: no augmentation
test_transform = transforms.Compose([
    transforms.Resize(256),                       # resize short side to 256
    transforms.CenterCrop(224),                   # center crop to 224
    transforms.ToTensor(),
    transforms.Normalize(IMAGENET_MEAN, IMAGENET_STD),
])

创建数据集

训练数据从文件夹中获取，在属性中设置训练数据的路径和预处理流水线

# dataset
train_data = datasets.ImageFolder(
    root='./data/train',
    transform=train_transform
)
val_data = datasets.ImageFolder(
    root='./data/val',
    transform=test_transform
)

加载模型，替换分类头

加载预训练模型，先冻结预训练模型中的原有参数，再将最后一层的输出变成2个

# Model 
model = models.resnet18(weights=None)             # no download
model.load_state_dict(
    torch.load("./cache/hub/checkpoints/resnet18-f37072fd.pth",
               map_location='cpu')
)

for param in model.parameters():
    param.requires_grad = False                   # freeze all

model.fc = nn.Linear(512, 2)        # input of fc layer is 512

训练＆保存

利用已经建立好的数据集加载数据，向前传播，反向传播，求导，求梯度，优化，循环数次后保存pth文件

if __name__ == '__main__':
    from torch.utils.data import DataLoader

    train_loader = DataLoader(train_data, batch_size=32, shuffle=True,
                          num_workers=2, pin_memory=True,
                          persistent_workers=True)
    val_loader = DataLoader(val_data, batch_size=32, shuffle=False,
                        num_workers=2, pin_memory=True,
                        persistent_workers=True)        # set to False if memory is tight

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model = model.to(device)

    criterion = nn.CrossEntropyLoss()
    optimizer = optim.SGD(
        filter(lambda p: p.requires_grad, model.parameters()),
        lr=0.001, momentum=0.9
    )

    for epoch in range(10):
        # Train
        model.train()
        running_loss = 0.0
        for imgs, labels in train_loader:
            imgs, labels = imgs.to(device), labels.to(device)

            optimizer.zero_grad()                 # clear gradients
            outputs = model(imgs)                 # forward
            loss = criterion(outputs, labels)     # compute loss
            loss.backward()                       # backward
            optimizer.step()                      # update
            running_loss += loss.item()

        print(f"Epoch {epoch+1}, Train Loss: {running_loss/len(train_loader):.4f}")

        # Validation
        model.eval()
        correct, total = 0, 0
        with torch.no_grad():
            for imgs, labels in val_loader:
                imgs, labels = imgs.to(device), labels.to(device)
                outputs = model(imgs)
                _, predicted = torch.max(outputs, 1)
                total += labels.size(0)
                correct += (predicted == labels).sum().item()

        print(f"Epoch {epoch+1}, Val Acc: {100 * correct / total:.2f}%")

    torch.save(model.state_dict(), "catdog_resnet18.pth")
    print("saved!")

Epoch 1, Train Loss: 0.1313
Epoch 1, Val Acc: 97.50%
Epoch 2, Train Loss: 0.0914
Epoch 2, Val Acc: 97.86%
Epoch 3, Train Loss: 0.0883
Epoch 3, Val Acc: 98.20%
Epoch 4, Train Loss: 0.0844
Epoch 4, Val Acc: 97.72%
Epoch 5, Train Loss: 0.0779
Epoch 5, Val Acc: 98.16%
Epoch 6, Train Loss: 0.0815
Epoch 6, Val Acc: 97.98%
Epoch 7, Train Loss: 0.0808
Epoch 7, Val Acc: 98.08%
Epoch 8, Train Loss: 0.0809
Epoch 8, Val Acc: 98.12%
Epoch 9, Train Loss: 0.0758
Epoch 9, Val Acc: 97.94%
Epoch 10, Train Loss: 0.0783
Epoch 10, Val Acc: 97.96%
saved!

可以尝试通过其他方法提高准确率

图形界面

插件：customtkinter

import os
import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision.models as models
from torchvision import transforms
from PIL import Image

import customtkinter as ctk
from tkinter import filedialog, messagebox


# ==================== Model ====================
# NOTE: must match training
def build_model():
    model = models.resnet18(weights=None)
    model.fc = nn.Linear(512, 2)
    return model


# ==================== Classes ====================
CLASS_NAMES = ["cat", "dog"]

DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = None
model_ready = False

# NOTE: must match test transform
IMAGENET_MEAN = [0.485, 0.456, 0.406]
IMAGENET_STD  = [0.229, 0.224, 0.225]

transform = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize(IMAGENET_MEAN, IMAGENET_STD),
])


def preprocess(pil_image):
    """RGB + resize 256 + center crop 224 + ToTensor + Normalize"""
    img = pil_image.convert("RGB")
    return transform(img).unsqueeze(0)


# ==================== Theme ====================
ctk.set_appearance_mode("dark")
ctk.set_default_color_theme("blue")

ACCENT = "#4da3ff"
CARD   = "#1a1a1a"
CARD_2 = "#222222"
TEXT   = "#ffffff"
SUB    = "#9a9a9a"
TRACK  = "#2a2a2a"

FONT       = ("Segoe UI", 12)
FONT_BOLD  = ("Segoe UI", 12, "bold")
FONT_TITLE = ("Segoe UI", 16, "bold")


# ==================== Animated bar ====================
class AnimatedBar(ctk.CTkProgressBar):
    def __init__(self, master, **kwargs):
        super().__init__(master, height=14,
                         progress_color=ACCENT,
                         fg_color=TRACK,
                         corner_radius=7, **kwargs)
        self.value = 0.0
        self.target = 0.0
        self._running = False
        self.set(0)

    def set_value(self, v):
        self.target = max(0.0, min(100.0, float(v))) / 100.0
        if not self._running:
            self._running = True
            self._animate()

    def _animate(self):
        diff = self.target - self.value
        if abs(diff) < 0.003:
            self.value = self.target
            self.set(self.value)
            self._running = False
            return
        self.value += diff * 0.25
        self.set(self.value)
        self.after(16, self._animate)


# ==================== App ====================
class App(ctk.CTk):
    def __init__(self):
        super().__init__()
        self.title("Cats vs Dogs Classifier")
        self.geometry("1000x680")
        self.minsize(880, 600)

        self.grid_columnconfigure(0, weight=1)
        self.grid_columnconfigure(1, weight=1)
        self.grid_rowconfigure(1, weight=1)

        self.external_img = None
        self._preview_tk = None

        self._build_topbar()
        self._build_left()
        self._build_right()

    # -------------------------------------------------
    def _build_topbar(self):
        bar = ctk.CTkFrame(self, fg_color=CARD, corner_radius=12)
        bar.grid(row=0, column=0, columnspan=2, sticky="ew",
                 padx=16, pady=(16, 0))
        bar.grid_columnconfigure(1, weight=1)

        ctk.CTkLabel(bar, text="Model", font=FONT_BOLD,
                     text_color=SUB).grid(row=0, column=0, padx=(14, 8), pady=10)
        self.model_path_var = ctk.StringVar(value="No model selected")
        ctk.CTkEntry(bar, textvariable=self.model_path_var,
                     state="readonly").grid(row=0, column=1, sticky="ew",
                                            padx=(0, 8), pady=10)
        ctk.CTkButton(bar, text="Load .pth", width=110, command=self.load_model,
                      fg_color=ACCENT, text_color="#0d0d0d",
                      hover_color="#6fd0ff", corner_radius=8)\
            .grid(row=0, column=2, padx=(0, 14), pady=10)

    # -------------------------------------------------
    def _build_left(self):
        left = ctk.CTkFrame(self, fg_color=CARD, corner_radius=12)
        left.grid(row=1, column=0, sticky="nsew", padx=16, pady=16)
        left.grid_rowconfigure(1, weight=1)
        left.grid_columnconfigure(0, weight=1)

        ctk.CTkLabel(left, text="Image", font=FONT_TITLE,
                     text_color=TEXT).grid(row=0, column=0, sticky="w",
                                           padx=18, pady=(14, 8))

        # preview area
        self.preview_label = ctk.CTkLabel(left, text="(no image)",
                                          fg_color="#111111",
                                          corner_radius=8,
                                          text_color=SUB)
        self.preview_label.grid(row=1, column=0, sticky="nsew",
                                padx=18, pady=(0, 14))

        btn_row = ctk.CTkFrame(left, fg_color="transparent")
        btn_row.grid(row=2, column=0, sticky="ew", padx=18, pady=(0, 14))

        ctk.CTkButton(btn_row, text="Load Image", width=120,
                      command=self.load_image,
                      fg_color=CARD_2, text_color=TEXT,
                      hover_color="#2e2e2e", corner_radius=8)\
            .pack(side="left")
        ctk.CTkButton(btn_row, text="Recognize", width=120,
                      command=self.recognize,
                      fg_color=ACCENT, text_color="#0d0d0d",
                      hover_color="#6fd0ff", corner_radius=8)\
            .pack(side="right")

    # -------------------------------------------------
    def _build_right(self):
        right = ctk.CTkFrame(self, fg_color=CARD, corner_radius=12)
        right.grid(row=1, column=1, sticky="nsew", padx=(0, 16), pady=16)
        right.grid_columnconfigure(0, weight=1)
        right.grid_rowconfigure(3, weight=1)

        ctk.CTkLabel(right, text="Result", font=FONT_TITLE,
                     text_color=TEXT).grid(row=0, column=0, sticky="w",
                                           padx=18, pady=(14, 8))

        self.result_var = ctk.StringVar(
            value="Load an image, then click Recognize")
        self.result_label = ctk.CTkLabel(
            right, textvariable=self.result_var,
            font=FONT_BOLD, text_color=TEXT,
            fg_color="#111111", corner_radius=8,
            anchor="w", justify="left",
            padx=14, pady=16,
        )
        self.result_label.grid(row=1, column=0, sticky="ew", padx=18)

        ctk.CTkLabel(right, text="Probabilities", font=FONT_BOLD,
                     text_color=SUB).grid(row=2, column=0, sticky="w",
                                          padx=18, pady=(18, 8))

        prob = ctk.CTkFrame(right, fg_color="transparent")
        prob.grid(row=3, column=0, sticky="nsew", padx=18, pady=(0, 14))
        prob.grid_columnconfigure(1, weight=1)

        self.bars = []
        self.bar_labels = []
        for i, name in enumerate(CLASS_NAMES):
            ctk.CTkLabel(prob, text=name, font=FONT, text_color=TEXT,
                         width=90, anchor="w")\
                .grid(row=i, column=0, sticky="w", pady=3)
            bar = AnimatedBar(prob)
            bar.grid(row=i, column=1, sticky="ew", padx=10, pady=3)
            val = ctk.CTkLabel(prob, text="0%", font=FONT, text_color=SUB,
                               width=50, anchor="e")
            val.grid(row=i, column=2, sticky="e", pady=3)
            self.bars.append(bar)
            self.bar_labels.append(val)

    # -------------------------------------------------
    def load_model(self):
        path = filedialog.askopenfilename(
            title="Select .pth model",
            filetypes=[("PyTorch model", "*.pth"), ("All files", "*.*")])
        if not path:
            return
        global model, model_ready
        try:
            m = build_model().to(DEVICE)
            m.load_state_dict(torch.load(path, map_location=DEVICE))
            m.eval()
            model = m
            model_ready = True
            self.model_path_var.set(path)
        except Exception as e:
            messagebox.showerror("Load failed", f"Failed to load model:\n{e}")

    # -------------------------------------------------
    def load_image(self):
        path = filedialog.askopenfilename(
            title="Select image",
            filetypes=[("Image", "*.png *.jpg *.jpeg *.bmp"),
                       ("All files", "*.*")])
        if not path:
            return
        self.external_img = Image.open(path)
        # fit preview to current label size
        self.update_idletasks()
        w = max(self.preview_label.winfo_width(), 200)
        h = max(self.preview_label.winfo_height(), 200)
        preview = self.external_img.convert("RGB").copy()
        preview.thumbnail((w - 16, h - 16), Image.BILINEAR)
        self._preview_tk = ctk.CTkImage(light_image=preview,
                                        dark_image=preview,
                                        size=preview.size)
        self.preview_label.configure(image=self._preview_tk, text="")

    # -------------------------------------------------
    def recognize(self):
        if not model_ready:
            messagebox.showwarning("Warning", "Please load a .pth model first")
            return
        if self.external_img is None:
            messagebox.showwarning("Warning", "Please load an image first")
            return

        tensor = preprocess(self.external_img).to(DEVICE)
        with torch.no_grad():
            out = model(tensor)
            prob = F.softmax(out, dim=1)[0]
            pred = out.argmax(dim=1).item()
            conf = prob[pred].item()

        self.result_var.set(
            f"Prediction: {CLASS_NAMES[pred]}    Confidence: {conf*100:.1f}%"
        )
        for i in range(2):
            p = prob[i].item() * 100
            self.bars[i].set_value(p)
            self.bar_labels[i].configure(text=f"{p:.0f}%")


if __name__ == "__main__":
    App().mainloop()