花卉数据集详解与准备指南

本教程采用公开的5类花卉数据集，涵盖雏菊、蒲公英、玫瑰、向日葵和郁金香。下载完成后，请将其解压至`flower_data`文件夹中。原始数据按类别分别存放于对应文件夹，结构如下所示：

分类类别：共包含 5 种花卉，对应 5 个文件夹： daisy（雏菊） dandelion（蒲公英） roses（玫瑰） sunflowers（向日葵） tulips（郁金香）

[image: 原始数据文件夹结构]

如果之前跑过YOLO等类似项目，应该清楚数据集的放置方式有一定规范。通常需要将数据划分为`train`和`val`两个文件夹，每个文件夹下再按类别创建子文件夹，并按照8:2的比例划分训练集和验证集。你可以借助AI编写的脚本自动整理，但务必仔细检查，避免数据错乱——这种低级错误一旦出现，排查起来会非常耗时。

[image: 整理后的数据文件夹结构]

整理完成后，只需将训练集和验证集的路径分别指向`train`和`val`文件夹即可。

从零开始编写训练代码

首先提供一个最基础的版本，确保整个流程能够顺利运行。下面给出完整的训练代码，其中包含了数据加载、模型定义、训练循环与验证逻辑，并加入了早停机制和学习率调度。

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
import os
import sys
from tqdm import tqdm

# 设置设备
device = torch.device('cuda' if torch.cuda.is_a vailable() else 'cpu')
print(f"使用设备: {device}")

# 数据路径
train_path = r"flower_data/train"
val_path = r"flower_data/val"

# 数据预处理
data_transform = {
    "train": transforms.Compose([
        transforms.RandomResizedCrop(224),
        transforms.RandomHorizontalFlip(),
        transforms.ToTensor(),
        transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
    ]),
    "val": transforms.Compose([
        transforms.Resize((224, 224)),
        transforms.ToTensor(),
        transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
    ])
}

# 加载数据
train_dataset = datasets.ImageFolder(train_path, transform=data_transform["train"])
val_dataset = datasets.ImageFolder(val_path, transform=data_transform["val"])
train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)
val_loader = DataLoader(val_dataset, batch_size=32, shuffle=False)
print(f"训练集大小: {len(train_dataset)}")
print(f"验证集大小: {len(val_dataset)}")
print(f"类别: {train_dataset.classes}")

# 定义AlexNet模型
class AlexNet(nn.Module):
    def __init__(self, num_classes=5):
        super(AlexNet, self).__init__()
        self.features = nn.Sequential(
            nn.Conv2d(3, 64, kernel_size=11, stride=4, padding=2), nn.ReLU(inplace=True), nn.MaxPool2d(kernel_size=3, stride=2),
            nn.Conv2d(64, 192, kernel_size=5, padding=2), nn.ReLU(inplace=True), nn.MaxPool2d(kernel_size=3, stride=2),
            nn.Conv2d(192, 384, kernel_size=3, padding=1), nn.ReLU(inplace=True),
            nn.Conv2d(384, 256, kernel_size=3, padding=1), nn.ReLU(inplace=True),
            nn.Conv2d(256, 256, kernel_size=3, padding=1), nn.ReLU(inplace=True), nn.MaxPool2d(kernel_size=3, stride=2),
        )
        self.classifier = nn.Sequential(
            nn.Dropout(),
            nn.Linear(256 * 6 * 6, 4096), nn.ReLU(inplace=True),
            nn.Dropout(),
            nn.Linear(4096, 4096), nn.ReLU(inplace=True),
            nn.Linear(4096, num_classes),
        )

    def forward(self, x):
        x = self.features(x)
        x = x.view(x.size(0), -1)
        x = self.classifier(x)
        return x

# 初始化模型、损失函数、优化器和调度器
model = AlexNet(num_classes=5).to(device)
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.0002)
scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=5, gamma=0.5)

# 训练函数
def train_epoch():
    model.train()
    total_loss = 0
    correct = 0
    total = 0
    train_bar = tqdm(train_loader, file=sys.stdout, desc="Training")
    for images, labels in train_bar:
        images, labels = images.to(device), labels.to(device)
        optimizer.zero_grad()
        outputs = model(images)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
        total_loss += loss.item()
        _, predicted = torch.max(outputs.data, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()
        train_bar.set_postfix({'loss': loss.item(), 'acc': correct / total})
    a vg_loss = total_loss / len(train_loader)
    accuracy = correct / total
    return a vg_loss, accuracy

# 验证函数
def validate():
    model.eval()
    correct = 0
    total = 0
    val_bar = tqdm(val_loader, file=sys.stdout, desc="Validating")
    with torch.no_grad():
        for images, labels in val_bar:
            images, labels = images.to(device), labels.to(device)
            outputs = model(images)
            _, predicted = torch.max(outputs.data, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()
            val_bar.set_postfix({'acc': correct / total})
    accuracy = correct / total
    return accuracy

# 训练模型
num_epochs = 50
best_acc = 0.0
patience = 10
patience_counter = 0
sa ve_path = r"alexnet_flower_optimized.pth"

print("\n开始训练...\n")
for epoch in range(num_epochs):
    train_loss, train_acc = train_epoch()
    val_acc = validate()
    scheduler.step()
    print(f"\nEpoch [{epoch+1}/{num_epochs}] | Loss: {train_loss:.4f} | Train Acc: {train_acc:.4f} | Val Acc: {val_acc:.4f}")

    if val_acc > best_acc:
        best_acc = val_acc
        torch.sa ve(model.state_dict(), sa ve_path)
        patience_counter = 0
        print(f" Best model sa ved! New best Val Acc: {best_acc:.4f}")
    else:
        patience_counter += 1
        print(f"Validation Acc did not improve. Patience: {patience_counter}/{patience}")

    if patience_counter >= patience:
        print(f" Early stopping triggered after {patience} epochs without improvement.")
        break

print('\nFinished Training')
print(f"最终最佳验证准确率: {best_acc:.4f}")

[image: 运行结果截图]

这里也分享一下，为什么在自己编写代码时很容易出现过拟合和欠拟合现象：

（1）一开始我对训练集和测试集使用了完全相同的预处理方式，结果导致了过拟合问题。

（2）对训练集采用了过于复杂的数据增强策略，反而引发了欠拟合问题。

这两种现象的根本原因，通常都出在数据处理环节。如果刚开始独立编写代码，建议参考成熟的代码库中数据预处理部分的写法，这样可以少走很多弯路。

逐步升级：项目模块化设计

下面我们将从以下几个模块来构建完整的项目：

（1）网络结构模块 （2）数据集读取模块 （3）训练文件模块 （4）测试文件模块 （5）辅助函数模块

[image: 项目结构图]

项目模块划分与功能概述

这种结构借鉴了YOLOv5的配置风格，能够清晰地将数据、模型和训练超参数分离。使用YAML配置文件进行深度学习项目管理，最大的优势在于实现了代码逻辑与配置参数的解耦，这是构建可复现项目的基础。

configs文件夹：配置文件管理

该结构模仿了YOLOv5的配置思路，能够清晰划分数据、模型和训练超参数。使用YAML配置文件进行深度学习项目管理，核心优势就是实现了代码逻辑与配置参数的清晰分离，这是构建专业、灵活且可复现的项目的基础。

flower.yaml中存放的内容：

数据路径 (Data Paths)：告诉程序训练和验证图片的存放位置。

类别信息 (Class Information)： `nc: 5`，`names: ['daisy', 'dandelion', 'roses', 'sunflowers', 'tulips']`

训练超参数 (Training Hyperparameters)： `epochs: 100`, `batch_size: 16`, `img_size: 224`, `lr0: 0.0001`, `weight_decay: 0.0001`

优化器和调度器 (Optimizer and Scheduler)： `optimizer: 'Adam'`, `scheduler: 'CosineAnnealing'`

模型参数 (Model Parameters)： `dropout: 0.5`

设备和路径 (Device and Paths)： `device: 'cuda:0'`, `workers: 4`, `sa ve_dir: 'runs/train'`

# 花卉分类配置文件
# 数据路径
train: 'flower_data/train'
val: 'flower_data/val'
# 类别信息
nc: 5
names: ['daisy', 'dandelion', 'roses', 'sunflowers', 'tulips']
# 训练超参数
epochs: 100
batch_size: 16
img_size: 224
lr0: 0.0001
weight_decay: 0.0001
# 优化器和调度器
optimizer: 'Adam'
scheduler: 'CosineAnnealing'
step_size: 10
gamma: 0.1
# 模型参数
dropout: 0.5
# 设备
device: 'cuda:0'
workers: 4
# 保存路径
sa ve_dir: 'runs/train'

[image: 配置文件示意图]

models文件夹：网络模型定义

网络结构文件（model文件）

models/Alexnet_model.py，用于定义网络的全部组件及前向传播逻辑。该模块主要存放网络的结构定义。

"""AlexNet模型"""
import torch.nn as nn
from models.common import ConvBNReLU, LinearBNReLU

class AlexNet(nn.Module):
    def __init__(self, num_classes=5, dropout=0.5):
        super().__init__()
        # 特征提取
        self.features = nn.Sequential(
            ConvBNReLU(3, 96, 11, 4, 2),
            nn.MaxPool2d(3, 2),
            ConvBNReLU(96, 256, 5, 1, 2),
            nn.MaxPool2d(3, 2),
            ConvBNReLU(256, 384, 3, 1, 1),
            ConvBNReLU(384, 384, 3, 1, 1),
            ConvBNReLU(384, 256, 3, 1, 1),
            nn.MaxPool2d(3, 2),
        )
        # 分类器
        self.classifier = nn.Sequential(
            nn.Dropout(dropout),
            LinearBNReLU(256 * 6 * 6, 2048),
            nn.Dropout(dropout),
            LinearBNReLU(2048, 1024),
            nn.Linear(1024, num_classes)
        )
        self._init_weights()
    
    def _init_weights(self):
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
            elif isinstance(m, (nn.BatchNorm2d, nn.BatchNorm1d)):
                nn.init.constant_(m.weight, 1)
                nn.init.constant_(m.bias, 0)
            elif isinstance(m, nn.Linear):
                nn.init.normal_(m.weight, 0, 0.01)
                nn.init.constant_(m.bias, 0)
    
    def forward(self, x):
        x = self.features(x)
        x = x.view(x.size(0), -1)
        x = self.classifier(x)
        return x

[image: 模型结构示意图]

models/common.py，用于定义网络中的通用组件，方便在Alexnet_model.py中直接调用。

"""通用模块"""
import torch.nn as nn

class ConvBNReLU(nn.Module):
    """卷积 + BN + ReLU"""
    def __init__(self, in_c, out_c, k=3, s=1, p=1):
        super().__init__()
        self.conv = nn.Conv2d(in_c, out_c, k, s, p, bias=False)
        self.bn = nn.BatchNorm2d(out_c)
        self.relu = nn.ReLU(inplace=True)
    
    def forward(self, x):
        return self.relu(self.bn(self.conv(x)))

class LinearBNReLU(nn.Module):
    """全连接 + BN + ReLU"""
    def __init__(self, in_f, out_f):
        super().__init__()
        self.fc = nn.Linear(in_f, out_f)
        self.bn = nn.BatchNorm1d(out_f)
        self.relu = nn.ReLU(inplace=True)
    
    def forward(self, x):
        return self.relu(self.bn(self.fc(x)))

[image: 通用模块示意图]

utils文件夹：工具函数与辅助模块

utils/datasets.py 数据集加载与预处理

"""数据集加载"""
from torch.utils.data import DataLoader
from torchvision import transforms, datasets

def create_dataloader(path, img_size=224, batch_size=16, shuffle=True, workers=4, augment=False):
    """创建数据加载器"""
    if augment:
        transform = transforms.Compose([
            transforms.Resize((img_size + 32, img_size + 32)),
            transforms.RandomCrop(img_size),
            transforms.RandomHorizontalFlip(),
            transforms.RandomRotation(15),
            transforms.ColorJitter(0.3, 0.3, 0.3, 0.1),
            transforms.ToTensor(),
            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
            transforms.RandomErasing(p=0.3)
        ])
    else:
        transform = transforms.Compose([
            transforms.Resize((img_size, img_size)),
            transforms.ToTensor(),
            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
        ])
    
    dataset = datasets.ImageFolder(path, transform=transform)
    loader = DataLoader(dataset, batch_size, shuffle, num_workers=workers, pin_memory=True)
    return loader, dataset

[image: 数据集加载示意图]

utils/general.py 通用工具函数

该模块包含`load_yaml`、`sa ve_checkpoint`、`load_checkpoint`三个核心函数，分别负责读取配置文件、保存模型检查点、加载模型检查点，是训练流程中不可或缺的“后勤保障”。

`load_yaml`函数的作用是读取一份训练项目的“配置清单”，并将其转换为Python可操作的字典格式。

`sa ve_checkpoint`函数的作用是为模型训练进度拍摄“快照”，并确保你能找到性能最优的版本。

`load_checkpoint`函数的作用是让一个中断的训练任务能够从上次停止的位置无缝接续，或者让一个训练好的模型准备好进行推理预测。

"""通用工具函数"""
import yaml
import torch
from pathlib import Path

def load_yaml(path):
    """加载YAML配置"""
    with open(path, 'r', encoding='utf-8') as f:
        return yaml.safe_load(f)

def sa ve_checkpoint(state, path, is_best=False):
    """保存检查点"""
    Path(path).parent.mkdir(parents=True, exist_ok=True)
    torch.sa ve(state, path)
    if is_best:
        best_path = Path(path).parent / 'best.pt'
        torch.sa ve(state, best_path)
        print(f' Best model sa ved: {best_path}')

def load_checkpoint(path, model, optimizer=None):
    """加载检查点"""
    ckpt = torch.load(path, map_location='cpu')
    model.load_state_dict(ckpt['model'])
    if optimizer and 'optimizer' in ckpt:
        optimizer.load_state_dict(ckpt['optimizer'])
    return ckpt.get('epoch', 0)

[image: 通用工具函数示意图]

utils/metrics.py 评估指标与统计工具

这里的`Metrics`类就像一个计分板，在每一轮训练（批次）中记录得分（预测和标签），直到整个Epoch结束，然后输出最终的比赛得分（各项指标）。`A verageMeter`类则是一个“在线平均值计算器”，可以高效地计算并随时提供当前的平均损失或准确率。

"""评估指标"""
import numpy as np
from sklearn.metrics import accuracy_score, precision_recall_fscore_support

class Metrics:
    """指标计算器"""
    def __init__(self):
        self.reset()
    
    def reset(self):
        self.preds = []
        self.labels = []
    
    def update(self, pred, label):
        self.preds.extend(pred.cpu().numpy())
        self.labels.extend(label.cpu().numpy())
    
    def compute(self):
        acc = accuracy_score(self.labels, self.preds)
        p, r, f1, _ = precision_recall_fscore_support(
            self.labels, self.preds, a verage='macro', zero_division=0
        )
        return {'acc': acc, 'precision': p, 'recall': r, 'f1': f1}

class A verageMeter:
    """平均值计算器"""
    def __init__(self):
        self.reset()
    
    def reset(self):
        self.val = 0
        self.a vg = 0
        self.sum = 0
        self.count = 0
    
    def update(self, val, n=1):
        self.val = val
        self.sum += val * n
        self.count += n
        self.a vg = self.sum / self.count

[image: 评估指标示意图]

utils/torch_utils.py PyTorch工具函数

"""PyTorch工具"""
import torch
import torch.nn as nn

def select_device(device=''):
    """选择设备"""
    if device and 'cuda' in device and torch.cuda.is_a vailable():
        return torch.device(device)
    return torch.device('cpu')

def get_optimizer(model, name='Adam', lr=0.001, weight_decay=0.0001):
    """创建优化器"""
    if name == 'SGD':
        return torch.optim.SGD(model.parameters(), lr, momentum=0.9, weight_decay=weight_decay)
    elif name == 'Adam':
        return torch.optim.Adam(model.parameters(), lr, weight_decay=weight_decay)
    else:
        raise ValueError(f'Unknown optimizer: {name}')

def get_scheduler(optimizer, name='CosineAnnealing', epochs=100, step_size=10, gamma=0.1):
    """创建学习率调度器"""
    if name == 'StepLR':
        return torch.optim.lr_scheduler.StepLR(optimizer, step_size, gamma)
    elif name == 'CosineAnnealing':
        return torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, epochs)
    return None

[image: 工具函数示意图]

主文件夹：核心训练与评估脚本

该目录包含`train.py`、`val.py`、`predict.py`三个核心脚本，以及`runs`和`weights`文件夹。

train.py：训练脚本

这是训练脚本，通过`argparse`接收命令行参数，然后加载配置、数据、模型，开始训练循环，并保存最佳模型。

"""训练脚本"""
import argparse
import torch
import torch.nn as nn
from pathlib import Path
from tqdm import tqdm
from models.AlexNet import AlexNet
from utils.datasets import create_dataloader
from utils.general import load_yaml, sa ve_checkpoint
from utils.metrics import Metrics, A verageMeter
from utils.torch_utils import select_device, get_optimizer, get_scheduler

def train_epoch(model, loader, criterion, optimizer, device):
    """训练一个epoch"""
    model.train()
    loss_meter = A verageMeter()
    metrics = Metrics()
    
    pbar = tqdm(loader, desc='Training')
    for imgs, labels in pbar:
        imgs, labels = imgs.to(device), labels.to(device)
        optimizer.zero_grad()
        outputs = model(imgs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
        loss_meter.update(loss.item(), imgs.size(0))
        metrics.update(outputs.argmax(1), labels)
        pbar.set_postfix({'loss': f'{loss_meter.a vg:.4f}'})
    
    results = metrics.compute()
    results['loss'] = loss_meter.a vg
    return results

def validate(model, loader, criterion, device):
    """验证"""
    model.eval()
    loss_meter = A verageMeter()
    metrics = Metrics()
    
    with torch.no_grad():
        for imgs, labels in tqdm(loader, desc='Validating'):
            imgs, labels = imgs.to(device), labels.to(device)
            outputs = model(imgs)
            loss = criterion(outputs, labels)
            loss_meter.update(loss.item(), imgs.size(0))
            metrics.update(outputs.argmax(1), labels)
    
    results = metrics.compute()
    results['loss'] = loss_meter.a vg
    return results

def main(opt):
    # 加载配置
    cfg = load_yaml(opt.cfg)
    device = select_device(opt.device or cfg['device'])
    sa ve_dir = Path(opt.sa ve_dir or cfg['sa ve_dir'])
    sa ve_dir.mkdir(parents=True, exist_ok=True)
    print(f' Training on {device}')
    
    # 数据加载
    train_loader, train_set = create_dataloader(
        cfg['train'], cfg['img_size'], cfg['batch_size'], 
        True, cfg['workers'], augment=True
    )
    val_loader, val_set = create_dataloader(
        cfg['val'], cfg['img_size'], cfg['batch_size'], 
        False, cfg['workers'], augment=False
    )
    print(f'Train: {len(train_set)}, Val: {len(val_set)}')
    
    # 模型
    model = AlexNet(cfg['nc'], cfg['dropout']).to(device)
    criterion = nn.CrossEntropyLoss(label_smoothing=0.1)
    optimizer = get_optimizer(model, cfg['optimizer'], cfg['lr0'], cfg['weight_decay'])
    scheduler = get_scheduler(optimizer, cfg['scheduler'], cfg['epochs'], 
                              cfg.get('step_size', 10), cfg.get('gamma', 0.1))
    
    # 训练
    best_acc = 0
    for epoch in range(1, cfg['epochs'] + 1):
        print(f'\n Epoch {epoch}/{cfg["epochs"]}')
        train_results = train_epoch(model, train_loader, criterion, optimizer, device)
        val_results = validate(model, val_loader, criterion, device)
        if scheduler:
            scheduler.step()
        print(f'Train - Loss: {train_results["loss"]:.4f}, Acc: {train_results["acc"]:.4f}')
        print(f'Val   - Loss: {val_results["loss"]:.4f}, Acc: {val_results["acc"]:.4f}')
        
        # 保存
        is_best = val_results['acc'] > best_acc
        if is_best:
            best_acc = val_results['acc']
        sa ve_checkpoint({
            'epoch': epoch,
            'model': model.state_dict(),
            'optimizer': optimizer.state_dict(),
            'best_acc': best_acc
        }, sa ve_dir / 'last.pt', is_best)
    
    print(f'\n Training complete! Best Acc: {best_acc:.4f}')

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--cfg', type=str, default='configs/flower.yaml', help='config file')
    parser.add_argument('--device', type=str, default='', help='cuda device, i.e. 0 or cpu')
    parser.add_argument('--sa ve-dir', type=str, default='', help='sa ve directory')
    opt = parser.parse_args()
    main(opt)

[image: train.py运行结果截图]

val.py：验证脚本

验证脚本，加载训练好的权重，在验证集上评估模型性能，并输出详细的指标，包括准确率、精确率、召回率、F1分数，以及混淆矩阵和分类报告。

"""验证脚本"""
import argparse
import torch
from tqdm import tqdm
from sklearn.metrics import confusion_matrix, classification_report
from models.AlexNet import AlexNet
from utils.datasets import create_dataloader
from utils.general import load_yaml, load_checkpoint
from utils.metrics import Metrics
from utils.torch_utils import select_device

def main(opt):
    # 加载配置
    cfg = load_yaml(opt.cfg)
    device = select_device(opt.device or cfg['device'])
    print(' Validation')
    
    # 加载模型
    model = AlexNet(cfg['nc'], cfg['dropout']).to(device)
    load_checkpoint(opt.weights, model)
    model.eval()
    
    # 加载数据
    val_loader, val_set = create_dataloader(
        cfg['val'], cfg['img_size'], batch_size=1, 
        shuffle=False, workers=0, augment=False
    )
    print(f'Dataset: {len(val_set)} images')
    
    # 验证
    metrics = Metrics()
    all_preds, all_labels = [], []
    
    with torch.no_grad():
        for imgs, labels in tqdm(val_loader):
            imgs = imgs.to(device)
            outputs = model(imgs)
            preds = outputs.argmax(1)
            metrics.update(preds, labels)
            all_preds.extend(preds.cpu().numpy())
            all_labels.extend(labels.numpy())
    
    # 结果
    results = metrics.compute()
    print(f'\n Results:')
    print(f'Accuracy:  {results["acc"]:.4f}')
    print(f'Precision: {results["precision"]:.4f}')
    print(f'Recall:    {results["recall"]:.4f}')
    print(f'F1-Score:  {results["f1"]:.4f}')
    print(f'\n Confusion Matrix:')
    print(confusion_matrix(all_labels, all_preds))
    print(f'\n Classification Report:')
    print(classification_report(all_labels, all_preds, target_names=cfg['names'], digits=4))

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--cfg', type=str, default='configs/flower.yaml', help='config file')
    parser.add_argument('--weights', type=str, default='runs/train/best.pt', help='model weights')
    parser.add_argument('--device', type=str, default='', help='cuda device')
    opt = parser.parse_args()
    main(opt)

[image: val.py运行结果截图]

predict.py：预测脚本

预测脚本，对单张图片进行推理，并输出Top-K的预测结果和置信度。

"""预测脚本"""
import argparse
import torch
from PIL import Image
from torchvision import transforms
from models.AlexNet import AlexNet
from utils.general import load_yaml, load_checkpoint
from utils.torch_utils import select_device

def main(opt):
    # 加载配置
    cfg = load_yaml(opt.cfg)
    device = select_device(opt.device or cfg['device'])
    print(' Prediction')
    
    # 加载模型
    model = AlexNet(cfg['nc'], cfg['dropout']).to(device)
    load_checkpoint(opt.weights, model)
    model.eval()
    
    # 图像预处理
    transform = transforms.Compose([
        transforms.Resize((cfg['img_size'], cfg['img_size'])),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ])
    
    img = Image.open(opt.source).convert('RGB')
    img_tensor = transform(img).unsqueeze(0).to(device)
    
    # 预测
    with torch.no_grad():
        outputs = model(img_tensor)
        probs = torch.softmax(outputs, dim=1)
        pred_class = probs.argmax(1).item()
        confidence = probs[0][pred_class].item()
        top_k_prob, top_k_idx = torch.topk(probs, opt.top_k)
    
    # 结果
    print(f'\n Prediction:')
    print(f'Class: {cfg["names"][pred_class]}')
    print(f'Confidence: {confidence:.4f} ({confidence*100:.2f}%)')
    print(f'\n Top-{opt.top_k}:')
    for i, (prob, idx) in enumerate(zip(top_k_prob[0], top_k_idx[0]), 1):
        name = cfg["names"][idx.item()]
        print(f'{i}. {name:12s} - {prob.item():.4f} ({prob.item()*100:.2f}%)')

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--source', type=str, required=True, help='image path')
    parser.add_argument('--cfg', type=str, default='configs/flower.yaml', help='config file')
    parser.add_argument('--weights', type=str, default='runs/train/best.pt', help='model weights')
    parser.add_argument('--device', type=str, default='', help='cuda device')
    parser.add_argument('--top-k', type=int, default=3, help='top k predictions')
    opt = parser.parse_args()
    main(opt)

[image: predict.py运行结果截图]

预测结果：从输出可以看到，模型预测的类别是正确的。

[image: 预测结果截图]