【AI 进阶笔记】cv 小白上手 Kaggle Bird vs Drone 【图像分类】

class BirdVsDroneSimpleDataset(Dataset):
    def __init__(self, img_dir, transform=None):
        self.img_dir = img_dir
        self.transform = transform
        self.img_paths = [os.path.join(img_dir, fname) for fname in os.listdir(img_dir) if fname.endswith('.jpg')]

    def __len__(self):
        return len(self.img_paths)

    def __getitem__(self, idx):
        img_path = self.img_paths[idx]
        img = Image.open(img_path).convert('RGB')

        # 用文件名判断标签：B 开头是鸟，D 开头是无人机
        fname = os.path.basename(img_path)
        if fname.startswith('B'):
            label = 0  # 鸟
        elif fname.startswith('D'):
            label = 1  # 无人机
        else:
            raise ValueError(f"Unknown file name format: {fname}")

        if self.transform:
            img = self.transform(img)

        return img, label

transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor()
])

train_dataset = BirdVsDroneSimpleDataset("train/images", transform=transform)
train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True, num_workers=4)

import torch.nn as nn
import torchvision.models as models

# 使用 ResNet50 预训练模型
model = models.resnet50(pretrained=True)

# 修改最后一层全连接层，以适应2分类问题
num_ftrs = model.fc.in_features
model.fc = nn.Linear(num_ftrs, 2)  # 2类（鸟 vs 无人机）

# 发送到 GPU（如果可用）
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = model.to(device)

import torch.optim as optim

criterion = nn.CrossEntropyLoss()
optimizer = optim.AdamW(model.parameters(), lr=1e-4)

def train_model(model, train_loader, val_loader, criterion, optimizer, num_epochs=10):
    train_losses, val_losses = [], []
    train_accuracies, val_accuracies = [], []

    for epoch in range(num_epochs):
        # ---------- 训练 ----------
        model.train()
        running_loss = 0.0
        correct, total = 0, 0
        for inputs, labels in train_loader:
            inputs, labels = inputs.to(device), labels.to(device)
            optimizer.zero_grad()
            outputs = model(inputs)
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()

            running_loss += loss.item() * inputs.size(0)
            preds = torch.argmax(outputs, dim=1)
            correct += (preds == labels).sum().item()
            total += labels.size(0)

        epoch_loss = running_loss / total
        epoch_acc = correct / total
        train_losses.append(epoch_loss)
        train_accuracies.append(epoch_acc)

        # ---------- 验证 ----------
        model.eval()
        val_loss = 0.0
        val_correct, val_total = 0, 0
        with torch.no_grad():
            for inputs, labels in val_loader:
                inputs, labels = inputs.to(device), labels.to(device)
                outputs = model(inputs)
                loss = criterion(outputs, labels)
                val_loss += loss.item() * inputs.size(0)

                preds = torch.argmax(outputs, dim=1)
                val_correct += (preds == labels).sum().item()
                val_total += labels.size(0)

        val_epoch_loss = val_loss / val_total
        val_epoch_acc = val_correct / val_total
        val_losses.append(val_epoch_loss)
        val_accuracies.append(val_epoch_acc)

        # ---------- 日志打印 ----------
        print(f"[{epoch+1}/{num_epochs}] Train Loss: {epoch_loss:.4f}, Acc: {epoch_acc:.4f} | Val Loss: {val_epoch_loss:.4f}, Acc: {val_epoch_acc:.4f}")

    return train_losses, train_accuracies, val_losses, val_accuracies

import matplotlib.pyplot as plt

def plot_metrics(train_losses, train_accuracies, val_losses, val_accuracies):
    epochs = range(1, len(train_losses) + 1)
    
    plt.figure(figsize=(12, 5))

    # Loss
    plt.subplot(1, 2, 1)
    plt.plot(epochs, train_losses, 'b', label='Train Loss')
    plt.plot(epochs, val_losses, 'r', label='Val Loss')
    plt.title('Loss over Epochs')
    plt.xlabel('Epoch')
    plt.ylabel('Loss')
    plt.legend()

    # Accuracy
    plt.subplot(1, 2, 2)
    plt.plot(epochs, train_accuracies, 'b', label='Train Acc')
    plt.plot(epochs, val_accuracies, 'r', label='Val Acc')
    plt.title('Accuracy over Epochs')
    plt.xlabel('Epoch')
    plt.ylabel('Accuracy')
    plt.legend()

    plt.tight_layout()
    plt.show()

from PIL import Image
import torch.nn.functional as F

def predict_image(model, image_path, transform, class_names=['bird', 'drone']):
    model.eval()
    image = Image.open(image_path).convert('RGB')
    img_tensor = transform(image).unsqueeze(0).to(device)

    with torch.no_grad():
        output = model(img_tensor)
        probs = F.softmax(output, dim=1)
        pred_class = torch.argmax(probs).item()

    print(f"Prediction: {class_names[pred_class]}, Confidence: {probs[0][pred_class]:.4f}")
    return image

image_path = "test/images/BT/BIRD_123.jpg"  
predict_image(model, image_path, transform)

from torchvision import transforms

# 数据增强（训练集）
train_transform = transforms.Compose([
    transforms.RandomResizedCrop(224),  # 随机裁剪并缩放到224x224
    transforms.RandomHorizontalFlip(),  # 随机水平翻转
    transforms.RandomRotation(20),  # 随机旋转 -20 到 20 度
    transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2, hue=0.2),  # 随机调整亮度、对比度等
    transforms.ToTensor(),  # 转换为Tensor
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),  # 正常化
])

# 验证集一般不需要增强，只做归一化
valid_transform = transforms.Compose([
    transforms.Resize(256),  # 调整为统一大小
    transforms.CenterCrop(224),  # 中心裁剪为224x224
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])

# 训练集应用增强
train_dataset = BirdVsDroneSimpleDataset("train/images", transform=train_transform)
train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True, num_workers=4)

# 验证集只做归一化
val_dataset = BirdVsDroneSimpleDataset("valid/images", transform=valid_transform)
val_loader = DataLoader(val_dataset, batch_size=32, shuffle=False, num_workers=4)

import torch.nn as nn
import torchvision.models as models

class ResNetWithDropout(nn.Module):
    def __init__(self, num_classes=2):
        super().__init__()
        self.base_model = models.resnet50(pretrained=True)
        in_features = self.base_model.fc.in_features
        self.base_model.fc = nn.Sequential(
            nn.Dropout(0.5),  # 添加Dropout
            nn.Linear(in_features, num_classes)
        )

    def forward(self, x):
        return self.base_model(x)

optimizer = torch.optim.Adam(model.parameters(), lr=1e-4, weight_decay=1e-4)

best_val_loss = float('inf')
patience = 3
counter = 0

for epoch in range(num_epochs):
    train(...)  # 训练逻辑
    val_loss = validate(...)

    if val_loss < best_val_loss:
        best_val_loss = val_loss
        counter = 0
        # 保存模型
        torch.save(model.state_dict(), "best_model.pth")
    else:
        counter += 1
        if counter >= patience:
            print("早停触发，终止训练")
            break

criterion = nn.CrossEntropyLoss(label_smoothing=0.1)

pip install efficientnet-pytorch

import torch
import torch.nn as nn
from efficientnet_pytorch import EfficientNet

class EfficientNetWithDropout(nn.Module):
    def __init__(self, num_classes=2):
        super().__init__()
        # 加载预训练的 EfficientNetB0（可以选择不同的版本：B0, B1, B2, B3, B4）
        self.base_model = EfficientNet.from_pretrained('efficientnet-b0')
        
        # 替换全连接层
        in_features = self.base_model._fc.in_features
        self.base_model._fc = nn.Sequential(
            nn.Dropout(0.5),  # Dropout
            nn.Linear(in_features, num_classes)  # 输出层，假设是二分类
        )

    def forward(self, x):
        return self.base_model(x)

model = EfficientNetWithDropout(num_classes=2).to(device)

from efficientnet_pytorch import EfficientNet
from torchvision.models import convnext_base, ConvNeXt_Base_Weights

class EnsembleModel(nn.Module):
    def __init__(self, num_classes=2):
        super().__init__()
        # EfficientNet
        self.efficientnet = EfficientNet.from_pretrained('efficientnet-b0')
        in_feat_e = self.efficientnet._fc.in_features
        self.efficientnet._fc = nn.Linear(in_feat_e, num_classes)

        # ConvNeXt
        self.convnext = convnext_base(weights=ConvNeXt_Base_Weights.IMAGENET1K_V1)
        in_feat_c = self.convnext.classifier[2].in_features
        self.convnext.classifier[2] = nn.Linear(in_feat_c, num_classes)

    def forward(self, x):
        out1 = self.efficientnet(x)
        out2 = self.convnext(x)

        # Soft Voting (平均两个输出的 softmax 概率)
        prob1 = torch.softmax(out1, dim=1)
        prob2 = torch.softmax(out2, dim=1)

        final_out = (prob1 + prob2) / 2
        return final_out

# 冻结特征提取层，只训练最后的线性层
for param in model.efficientnet.parameters():
    param.requires_grad = False
for param in model.convnext.parameters():
    param.requires_grad = False

# 只解冻最后一层
for param in model.efficientnet._fc.parameters():
    param.requires_grad = True
for param in model.convnext.classifier[2].parameters():
    param.requires_grad = True

criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(filter(lambda p: p.requires_grad, model.parameters()), lr=1e-4)

final_out = 0.7 * prob1 + 0.3 * prob2  # 给 EfficientNet 更高权重

learning_rates = [1e-3, 5e-4, 1e-4]
for lr in learning_rates:
    optimizer = torch.optim.Adam(model.parameters(), lr=lr)
    train_model(model, ...)

from itertools import product

lrs = [1e-3, 1e-4]
bss = [32, 64]
wds = [0.0, 1e-4]

for lr, bs, wd in product(lrs, bss, wds):
    print(f"正在尝试 lr={lr}, bs={bs}, weight_decay={wd}")
    # 设置 optimizer、DataLoader 等重新训练

import optuna

def objective(trial):
    lr = trial.suggest_loguniform('lr', 1e-5, 1e-2)
    optimizer = torch.optim.Adam(model.parameters(), lr=lr)
    acc = train_and_eval(model, optimizer)
    return acc  # 越高越好

study = optuna.create_study(direction='maximize')
study.optimize(objective, n_trials=20)
print("Best params:", study.best_params)