基于PyTorch的自主编写VGG16猫狗分类实验代码如下：
# 1. 导入必要库
import torch
import torch.nn as nn
import torch.optim as optim
from torch.optim.lr_scheduler import StepLR
from torchvision import transforms
from torch.utils.data import DataLoader, ConcatDataset
import os
from PIL import Image
import matplotlib.pyplot as plt
from sklearn.metrics import accuracy_score, precision_score, recall_score, confusion_matrix
import numpy as np
# 2. 超参数与设备配置
batch_size = 16 # GPU内存不足时改为8/4，CPU改为4
epochs = 10
lr = 1e-4
gamma = 0.7
num_classes = 2 # 猫/狗二分类
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
root = r"d:\Users\86192\Desktop\catdog\catdog" # 数据集根路径
# 3. 数据预处理（训练集增强，测试/验证集基础处理）
train_transform = transforms.Compose([
transforms.RandomResizedCrop(224),
transforms.RandomHorizontalFlip(p=0.5),
transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2, hue=0.1),
transforms.GaussianBlur(kernel_size=(3, 3), sigma=(0.1, 2.0)),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])
val_test_transform = transforms.Compose([
transforms.Resize(256),
transforms.CenterCrop(224),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])
# 自定义数据集：从文件名提取猫/狗标签（文件名含"cat"→猫，含"dog"→狗）
class CustomDataset(torch.utils.data.Dataset):
def __init__(self, img_dir, transform=None):
self.img_dir = img_dir
self.transform = transform
# 筛选图片文件（支持.jpg/.jpeg/.png）
self.img_paths = [
os.path.join(img_dir, f)
for f in os.listdir(img_dir)
if f.endswith(('.jpg', '.jpeg', '.png'))
]
# 从文件名提取标签（0=猫，1=狗）
self.labels = []
for path in self.img_paths:
filename = os.path.basename(path).lower()
if 'cat' in filename:
self.labels.append(0)
elif 'dog' in filename:
self.labels.append(1)
else:
raise ValueError(f"文件名 {path} 不含'cat'/'dog'，无法自动标注！请检查文件名或手动标注。")
def __len__(self):
return len(self.img_paths)
def __getitem__(self, idx):
img = Image.open(self.img_paths[idx]).convert('RGB') # 统一转RGB（避免灰度图报错）
label = self.labels[idx]
if self.transform:
img = self.transform(img)
return img, label
# 4. 加载数据集（使用CustomDataset替代ImageFolder）
# 训练集：合并training_set + training_set2
train_dirs = [os.path.join(root, "training_set"), os.path.join(root, "training_set2")]
train_datasets = [CustomDataset(dir, transform=train_transform) for dir in train_dirs]
train_dataset = ConcatDataset(train_datasets)
train_loader = DataLoader(
train_dataset, batch_size=batch_size, shuffle=True, num_workers=2
)
# 验证集：training_set3
val_dataset = CustomDataset(os.path.join(root, "training_set3"), transform=val_test_transform)
val_loader = DataLoader(
val_dataset, batch_size=batch_size, shuffle=False, num_workers=2
)
# 测试集：test_set
test_dataset = CustomDataset(os.path.join(root, "test_set"), transform=val_test_transform)
test_loader = DataLoader(
test_dataset, batch_size=batch_size, shuffle=False, num_workers=2
)
# 小测试集：test_set3
test_small_dataset = CustomDataset(os.path.join(root, "test_set3"), transform=val_test_transform)
test_small_loader = DataLoader(
test_small_dataset, batch_size=batch_size, shuffle=False, num_workers=2
)
# 打印数据集统计（验证标签提取是否正确）
print("=== 数据集统计 ===")
for ds, name in zip(
[train_datasets[0], train_datasets[1], val_dataset, test_dataset, test_small_dataset],
["training_set", "training_set2", "training_set3", "test_set", "test_set3"]
):
cat_num = sum(1 for l in ds.labels if l == 0)
dog_num = len(ds.labels) - cat_num
print(f"{name}: 总图片数={len(ds)}, 猫={cat_num}, 狗={dog_num}")
# 5. 自主编写VGG16模型（完全复现实验报告结构）
class VGG16Custom(nn.Module):
def __init__(self, num_classes=2):
super(VGG16Custom, self).__init__()
# 特征提取层（13个卷积 + 5个池化，严格匹配VGG16结构）
self.features = nn.Sequential(
# 第1组：2×Conv3-64 + MaxPool
nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1),
nn.ReLU(inplace=True),
nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=1),
nn.ReLU(inplace=True),
nn.MaxPool2d(kernel_size=2, stride=2),
# 第2组：2×Conv3-128 + MaxPool
nn.Conv2d(64, 128, kernel_size=3, stride=1, padding=1),
nn.ReLU(inplace=True),
nn.Conv2d(128, 128, kernel_size=3, stride=1, padding=1),
nn.ReLU(inplace=True),
nn.MaxPool2d(kernel_size=2, stride=2),
# 第3组：3×Conv3-256 + MaxPool
nn.Conv2d(128, 256, kernel_size=3, stride=1, padding=1),
nn.ReLU(inplace=True),
nn.Conv2d(256, 256, kernel_size=3, stride=1, padding=1),
nn.ReLU(inplace=True),
nn.Conv2d(256, 256, kernel_size=3, stride=1, padding=1),
nn.ReLU(inplace=True),
nn.MaxPool2d(kernel_size=2, stride=2),
# 第4组：3×Conv3-512 + MaxPool
nn.Conv2d(256, 512, kernel_size=3, stride=1, padding=1),
nn.ReLU(inplace=True),
nn.Conv2d(512, 512, kernel_size=3, stride=1, padding=1),
nn.ReLU(inplace=True),
nn.Conv2d(512, 512, kernel_size=3, stride=1, padding=1),
nn.ReLU(inplace=True),
nn.MaxPool2d(kernel_size=2, stride=2),
# 第5组：3×Conv3-512 + MaxPool
nn.Conv2d(512, 512, kernel_size=3, stride=1, padding=1),
nn.ReLU(inplace=True),
nn.Conv2d(512, 512, kernel_size=3, stride=1, padding=1),
nn.ReLU(inplace=True),
nn.Conv2d(512, 512, kernel_size=3, stride=1, padding=1),
nn.ReLU(inplace=True),
nn.MaxPool2d(kernel_size=2, stride=2)
)
# 平均池化层（适配全连接层输入）
self.avgpool = nn.AdaptiveAvgPool2d((7, 7))
# 分类器（3个全连接层 + Dropout正则化）
self.classifier = nn.Sequential(
nn.Linear(512 * 7 * 7, 4096),
nn.ReLU(inplace=True),
nn.Dropout(p=0.5),
nn.Linear(4096, 4096),
nn.ReLU(inplace=True),
nn.Dropout(p=0.5),
nn.Linear(4096, num_classes)
)
# 实验要求：冻结所有参数，仅解冻features[0]层（第一个卷积层）
for param in self.parameters():
param.requires_grad = False
for param in self.features[0].parameters():
param.requires_grad = True
def forward(self, x):
x = self.features(x) # 特征提取
x = self.avgpool(x) # 池化
x = torch.flatten(x, 1) # 展平（batch维度外的维度合并）
x = self.classifier(x) # 分类
return x
# 实例化模型并转移到设备（GPU/CPU）
model = VGG16Custom(num_classes=num_classes).to(device)
print(f"模型已加载至{device}，仅解冻features[0]层（第一个卷积层）")
# 6. 损失函数、优化器、学习率调度器
criterion = nn.CrossEntropyLoss() # 交叉熵损失（适配二分类）
# 仅优化解冻的参数（features[0]层 + 分类器层）
optimizer = optim.Adam(
filter(lambda p: p.requires_grad, model.parameters()),
lr=lr
)
scheduler = StepLR(optimizer, step_size=3, gamma=gamma) # 每3轮学习率衰减为70%
# 7. 训练函数
def train(model, device, train_loader, criterion, optimizer, epoch):
model.train() # 启用训练模式（Dropout生效）
train_loss = 0.0
correct = 0
total = 0
for batch_idx, (data, target) in enumerate(train_loader):
data, target = data.to(device), target.to(device)
# 前向传播
optimizer.zero_grad() # 清空梯度
output = model(data)
loss = criterion(output, target)
# 反向传播 + 参数更新
loss.backward()
optimizer.step()
# 统计损失与准确率
train_loss += loss.item() * data.size(0)
_, predicted = output.max(1)
total += target.size(0)
correct += predicted.eq(target).sum().item()
# 打印批次进度
if batch_idx % 10 == 0:
progress = 100. * batch_idx / len(train_loader)
print(f'Train Epoch: {epoch} [{batch_idx * len(data)}/{len(train_loader.dataset)} '
f'({progress:.0f}%)]\tLoss: {loss.item():.6f}')
# 计算平均损失与准确率
avg_loss = train_loss / len(train_loader.dataset)
acc = 100. * correct / total
print(f'Epoch {epoch} | 训练损失: {avg_loss:.6f} | 训练准确率: {acc:.2f}%\n')
return avg_loss, acc
# 8. 验证函数
def validate(model, device, val_loader, criterion):
model.eval() # 启用评估模式（Dropout关闭）
val_loss = 0.0
correct = 0
total = 0
with torch.no_grad(): # 禁用梯度计算（节省内存）
for data, target in val_loader:
data, target = data.to(device), target.to(device)
output = model(data)
loss = criterion(output, target)
val_loss += loss.item() * data.size(0)
_, predicted = output.max(1)
total += target.size(0)
correct += predicted.eq(target).sum().item()
# 计算验证损失与准确率
avg_loss = val_loss / len(val_loader.dataset)
acc = 100. * correct / total
print(f'验证损失: {avg_loss:.6f} | 验证准确率: {acc:.2f}%\n')
return avg_loss, acc
# 9. 测试函数（输出Accuracy、Precision、Recall）
def test(model, device, test_loader, test_name):
model.eval()
all_preds = [] # 存储所有预测结果
all_targets = [] # 存储所有真实标签
with torch.no_grad():
for data, target in test_loader:
data, target = data.to(device), target.to(device)
output = model(data)
_, predicted = output.max(1)
# 转移到CPU并转为numpy（适配sklearn指标计算）
all_preds.extend(predicted.cpu().numpy())
all_targets.extend(target.cpu().numpy())
# 计算三大指标（正类=狗（1），负类=猫（0））
accuracy = accuracy_score(all_targets, all_preds)
precision = precision_score(
all_targets, all_preds,
average='binary', zero_division=0 # zero_division处理无正类预测的情况
)
recall = recall_score(
all_targets, all_preds,
average='binary', zero_division=0
)
# 计算混淆矩阵（辅助分析错误类型）
cm = confusion_matrix(all_targets, all_preds)
# 打印结果
print(f"=== {test_name} 评价指标 ===")
print(f"Accuracy（准确率）: {accuracy:.4f} ({accuracy*100:.2f}%)")
print(f"Precision（精确率）: {precision:.4f}")
print(f"Recall（召回率）: {recall:.4f}")
print(f"混淆矩阵:\n{cm}")
print(f"TP（狗预测为狗）: {cm[1,1]}, TN（猫预测为猫）: {cm[0,0]}")
print(f"FP（猫预测为狗）: {cm[0,1]}, FN（狗预测为猫）: {cm[1,0]}\n")
return accuracy, precision, recall, cm
# 10. 执行训练与测试
if __name__ == "__main__":
# 记录训练过程（用于绘图）
train_loss_history = []
train_acc_history = []
val_loss_history = []
val_acc_history = []
print("===== 开始训练 =====")
for epoch in range(1, epochs + 1):
train_loss, train_acc = train(model, device, train_loader, criterion, optimizer, epoch)
val_loss, val_acc = validate(model, device, val_loader, criterion)
scheduler.step() # 更新学习率
# 保存训练历史
train_loss_history.append(train_loss)
train_acc_history.append(train_acc)
val_loss_history.append(val_loss)
val_acc_history.append(val_acc)
print("===== 开始测试 =====")
# 测试主测试集（test_set）
test(model, device, test_loader, "主测试集（test_set）")
# 测试小测试集（test_set3）
test(model, device, test_small_loader, "小测试集（test_set3）")
# 绘制训练/验证曲线（保存到数据集根目录）
plt.figure(figsize=(12, 5))
# 子图1：损失曲线
plt.subplot(1, 2, 1)
plt.plot(range(1, epochs+1), train_loss_history, label="训练损失", color="blue")
plt.plot(range(1, epochs+1), val_loss_history, label="验证损失", color="red")
plt.xlabel("轮次（Epoch）")
plt.ylabel("损失（Loss）")
plt.title("VGG16训练与验证损失曲线")
plt.legend()
plt.grid(alpha=0.3)
# 子图2：准确率曲线
plt.subplot(1, 2, 2)
plt.plot(range(1, epochs+1), train_acc_history, label="训练准确率", color="blue")
plt.plot(range(1, epochs+1), val_acc_history, label="验证准确率", color="red")
plt.xlabel("轮次（Epoch）")
plt.ylabel("准确率（%）")
plt.title("VGG16训练与验证准确率曲线")
plt.legend()
plt.grid(alpha=0.3)
# 保存曲线
curve_path = os.path.join(root, "training_validation_curves.jpg")
plt.tight_layout()
plt.savefig(curve_path, dpi=300, bbox_inches="tight")
plt.show()
print(f"训练/验证曲线已保存至: {curve_path}")
# 保存模型参数
model_path = os.path.join(root, "vgg16_custom_catdog_model.pth")
torch.save(model.state_dict(), model_path)
print(f"自主编写VGG16模型参数已保存至: {model_path}")
代码是否存在错误，是否有更优化的代码或者是更符合VGG16模型的代码，请给出我完整的代码
扫一扫
5556
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
