数据集介绍
花分类数据集：
百度云链接下载: https://pan.baidu.com/s/1QLCTA4sXnQAw_yvxPj9szg
提取码:58p0
下载好之后，解压到flower_data文件夹下，此时flower_data\flower_photos下就是放的我们的数据
集，我们看一下原始的数据是什么样子的：
分类类别：共包含 5 类花卉，对应 5 个文件夹： daisy（雏菊） dandelion（蒲公英） roses（玫
瑰） sunflowers（向日葵） tulips（郁金香）
跑过一些项目的应该都有印象，比如YOLO等，他们的数据集的放置是有要求的一般情况下都是分
成两个，一个是train文件夹，train文件夹下是各种分类的文件夹（每个文件夹的名字是类报名）。
另外一个是val文件夹，val文件夹下是各种分类的文件夹（每个文件夹的名字是类报名）。一般是
按照8：2的比例去分这两个数据集的。这里的话可以用AI写代码整理，但是别忘记了检查一下。
整理好之后：
训练集的路径：D:\vscode\shenduxvexishizhan\VIT\flower_data\train
验证集的路径是：D:\vscode\shenduxvexishizhan\VIT\flower_data\val
打怪升级路线
下面我们将会从下面几个部分来做这个项目：
（1）网络结构模块
（2）数据集读取模块
（3）训练文件模块
（4）测试文件模块
（5）辅助函数模块
网络结构模块
这个模块在前面的文章已经详细介绍了：PyTorch深度学习实战 |手算ViT(Vision Transformer)模型-CSDN博客
假设我们的最后是1000个分类的：
假设原始图像的大小是224*224*3的
输入端：
【1】经过卷积，然后拉平之后变成196×768大小的张量
【2】经过Class Embedding，变成了197×768大小的张量
【3】经过Position Embedding，变成了197×768大小的张量
编码器：
【1】经过编码器之后，大小不发生变化，仍然是197×768大小的张量
【2】提取CLS Token，大小是1×768
输出端：
【1】经过一个线性层，变成了1×1000的张量，每个数字代表着这个类别的概率
代码实现：
import torch import torch.nn as nn import torch.nn.functional as F # 极简ViT实现（仅保留核心逻辑） class ViT(nn.Module): def __init__(self, img_size=224, patch_size=16, num_classes=5, embed_dim=384, num_heads=6, num_layers=6): super().__init__() # 1. Patch Embedding self.num_patches = (img_size // patch_size) ** 2 self.patch_embed = nn.Conv2d(3, embed_dim, kernel_size=patch_size, stride=patch_size) # 2. CLS Token + 位置嵌入 self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim)) self.pos_embed = nn.Parameter(torch.zeros(1, self.num_patches + 1, embed_dim)) # 3. Transformer编码器（简化版） self.transformer = nn.Sequential(*[ nn.TransformerEncoderLayer( d_model=embed_dim, nhead=num_heads, dim_feedforward=embed_dim*4, activation='gelu', batch_first=True ) for _ in range(num_layers) ]) # 4. 分类头 self.norm = nn.LayerNorm(embed_dim) self.head = nn.Linear(embed_dim, num_classes) def forward(self, x): # Patch嵌入: [B,3,224,224] → [B,384,14,14] → [B,196,384] x = self.patch_embed(x).flatten(2).transpose(1, 2) # 添加CLS Token: [B,197,384] cls_token = self.cls_token.expand(x.shape[0], -1, -1) x = torch.cat((cls_token, x), dim=1) # 添加位置嵌入 x = x + self.pos_embed # Transformer编码 x = self.transformer(x) # 分类（仅用CLS Token） x = self.norm(x)[:, 0] x = self.head(x) return x # 快速创建模型 def create_vit_model(): return ViT()
这段代码应该没有什么好说的了，模块最后的输出应该是batch*numclass的张量，其中batch，每
个批次的图片的数量，numclass表示预测的类别。批量花卉图片张量 [B,3,224,224] 分类预测得分 [B,5]
数据集读取模块
代码实现
import torch from torch.utils.data import DataLoader from torchvision import datasets, transforms # 数据集路径配置（替换为你的实际路径） TRAIN_PATH = r"D:\vscode\shenduxvexishizhan\deep-learning-for-image-processing-master\flower_data\train" VAL_PATH = r"D:\vscode\shenduxvexishizhan\deep-learning-for-image-processing-master\flower_data\val" def get_data_loaders(batch_size=16): """获取训练/验证数据加载器""" # 训练集预处理（含简单增强） train_transform = transforms.Compose([ transforms.Resize((224, 224)), transforms.RandomHorizontalFlip(), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ]) # 验证/测试集预处理 val_transform = transforms.Compose([ transforms.Resize((224, 224)), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ]) # 加载数据集 train_dataset = datasets.ImageFolder(TRAIN_PATH, transform=train_transform) val_dataset = datasets.ImageFolder(VAL_PATH, transform=val_transform) # 创建数据加载器 train_loader = DataLoader( train_dataset, batch_size=batch_size, shuffle=True, num_workers=0 # Windows下设为0 ) val_loader = DataLoader( val_dataset, batch_size=batch_size, shuffle=False, num_workers=0 ) # 类别映射（用于推理） class2idx = train_dataset.class_to_idx idx2class = {v: k for k, v in class2idx.items()} return train_loader, val_loader, idx2class
训练文件模块
代码实现
import torch import torch.nn as nn import torch.optim as optim from tqdm import tqdm from model import create_vit_model from dataset import get_data_loaders from utils import get_device, plot_curve, save_model # 训练配置 BATCH_SIZE = 16 EPOCHS = 20 LR = 1e-4 WEIGHT_DECAY = 1e-4 def train_one_epoch(model, loader, loss_fn, optimizer, device): """单轮训练""" model.train() total_loss, total_acc = 0, 0 pbar = tqdm(loader, desc='Training') for imgs, labels in pbar: imgs, labels = imgs.to(device), labels.to(device) # 前向传播 outputs = model(imgs) loss = loss_fn(outputs, labels) # 反向传播 optimizer.zero_grad() loss.backward() optimizer.step() # 统计指标 acc = (outputs.argmax(1) == labels).float().mean() total_loss += loss.item() total_acc += acc.item() # 更新进度条 pbar.set_postfix({'loss': loss.item(), 'acc': acc.item()}) return total_loss/len(loader), total_acc/len(loader) @torch.no_grad() def validate(model, loader, loss_fn, device): """单轮验证""" model.eval() total_loss, total_acc = 0, 0 pbar = tqdm(loader, desc='Validating') for imgs, labels in pbar: imgs, labels = imgs.to(device), labels.to(device) outputs = model(imgs) loss = loss_fn(outputs, labels) acc = (outputs.argmax(1) == labels).float().mean() total_loss += loss.item() total_acc += acc.item() pbar.set_postfix({'loss': loss.item(), 'acc': acc.item()}) return total_loss/len(loader), total_acc/len(loader) def main(): # 1. 初始化 device = get_device() print(f'Using device: {device}') # 2. 加载数据 train_loader, val_loader, _ = get_data_loaders(BATCH_SIZE) # 3. 创建模型 model = create_vit_model().to(device) loss_fn = nn.CrossEntropyLoss() optimizer = optim.AdamW(model.parameters(), lr=LR, weight_decay=WEIGHT_DECAY) scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=10, gamma=0.5) # 4. 训练记录 train_losses, train_accs = [], [] val_losses, val_accs = [], [] best_acc = 0 # 5. 训练循环 for epoch in range(EPOCHS): print(f'\nEpoch {epoch+1}/{EPOCHS}') # 训练 train_loss, train_acc = train_one_epoch(model, train_loader, loss_fn, optimizer, device) # 验证 val_loss, val_acc = validate(model, val_loader, loss_fn, device) # 学习率调度 scheduler.step() # 记录 train_losses.append(train_loss) train_accs.append(train_acc) val_losses.append(val_loss) val_accs.append(val_acc) # 保存最优模型 if val_acc > best_acc: best_acc = val_acc save_model(model) print(f'Best model saved! Val Acc: {best_acc:.4f}') # 打印本轮结果 print(f'Train: Loss={train_loss:.4f}, Acc={train_acc:.4f}') print(f'Val: Loss={val_loss:.4f}, Acc={val_acc:.4f}') # 6. 绘制训练曲线 plot_curve(train_losses, train_accs, val_losses, val_accs) print(f'\nTraining finished! Best Val Acc: {best_acc:.4f}') if __name__ == '__main__': main()
测试代码模块
代码实现
import torch from PIL import Image import torchvision.transforms as transforms import matplotlib.pyplot as plt from model import create_vit_model from dataset import VAL_PATH from utils import get_device, load_model # 测试配置 MODEL_PATH = 'best_model.pth' TEST_IMG_PATH = r'D:\vscode\shenduxvexishizhan\deep-learning-for-image-processing-master\flower_data\val\roses\123.jpg' # 替换为你的测试图片 def predict_single_image(): """单张图片测试/推理""" # 1. 初始化 device = get_device() model = create_vit_model().to(device) model = load_model(model, MODEL_PATH, device) model.eval() # 2. 获取类别映射 from torchvision.datasets import ImageFolder val_dataset = ImageFolder(VAL_PATH) idx2class = {v: k for k, v in val_dataset.class_to_idx.items()} # 3. 图片预处理 transform = transforms.Compose([ transforms.Resize((224, 224)), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ]) # 4. 加载并处理图片 img = Image.open(TEST_IMG_PATH).convert('RGB') img_tensor = transform(img).unsqueeze(0).to(device) # 5. 推理 with torch.no_grad(): output = model(img_tensor) prob = torch.softmax(output, dim=1) pred_idx = output.argmax(1).item() pred_class = idx2class[pred_idx] pred_conf = prob[0][pred_idx].item() # 6. 可视化结果 plt.imshow(img) plt.title(f'Prediction: {pred_class}\nConfidence: {pred_conf:.4f}') plt.axis('off') plt.show() # 7. 打印结果 print('='*30) print(f'Test Image: {TEST_IMG_PATH}') print(f'Predicted Class: {pred_class}') print(f'Confidence: {pred_conf:.4f}') print('='*30) # 打印所有类别概率 print('\nClass Probabilities:') for idx, cls in idx2class.items(): print(f'{cls}: {prob[0][idx]:.4f}') if __name__ == '__main__': predict_single_image()
辅助函数模块
代码实现
import torch import matplotlib.pyplot as plt # 设备选择 def get_device(): return torch.device("cuda" if torch.cuda.is_available() else "cpu") # 绘制训练曲线 def plot_curve(train_losses, train_accs, val_losses, val_accs): plt.figure(figsize=(12, 5)) # 损失曲线 plt.subplot(1,2,1) plt.plot(train_losses, label='Train Loss') plt.plot(val_losses, label='Val Loss') plt.xlabel('Epoch') plt.ylabel('Loss') plt.legend() plt.title('Loss Curve') # 准确率曲线 plt.subplot(1,2,2) plt.plot(train_accs, label='Train Acc') plt.plot(val_accs, label='Val Acc') plt.xlabel('Epoch') plt.ylabel('Accuracy') plt.legend() plt.title('Accuracy Curve') plt.tight_layout() plt.savefig('training_curve.png') plt.show() # 保存模型 def save_model(model, path='best_model.pth'): torch.save(model.state_dict(), path) # 加载模型 def load_model(model, path='best_model.pth', device='cuda'): model.load_state_dict(torch.load(path, map_location=device)) return model
实验结果
显然是过拟合了，后面的教程中，我们再去学习如何解决这个问题
