import os
import pandas as pd
import numpy as np
from PIL import Image
from tqdm import tqdm
from sklearn.model_selection import StratifiedKFold

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader, Subset
from torchvision import transforms, models
from torch.cuda.amp import GradScaler, autocast
from torch.optim.lr_scheduler import CosineAnnealingLR, ReduceLROnPlateau, OneCycleLR
from sklearn.model_selection import StratifiedKFold

# 定义 DogBreedDataset 类
class DogBreedDataset(Dataset):
    def __init__(self, csv_file, root_dir, transform=None):
        self.labels_df = pd.read_csv(csv_file)
        self.root_dir = root_dir
        self.transform = transform
        self.breeds = self.labels_df['breed'].unique()
        self.breed_to_idx = {breed: idx for idx, breed in enumerate(self.breeds)}
        self.labels_df['breed'] = self.labels_df['breed'].map(self.breed_to_idx)

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

    def __getitem__(self, idx):
        img_name = os.path.join(self.root_dir, self.labels_df.iloc[idx, 0] + '.jpg')
        image = Image.open(img_name).convert('RGB')
        breed = self.labels_df.iloc[idx, 1]
        if self.transform:
            image = self.transform(image)
        return image, breed

    def get_breed_name(self, idx):
        return self.breeds[idx]

# TestDataset 类定义
class TestDataset(Dataset):
    def __init__(self, root_dir, transform=None):
        self.root_dir = root_dir
        self.transform = transform
        self.image_ids = os.listdir(root_dir)

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

    def __getitem__(self, idx):
        img_name = os.path.join(self.root_dir, self.image_ids[idx])
        image = Image.open(img_name).convert('RGB')
        if self.transform:
            image = self.transform(image)
        return image, self.image_ids[idx].split('.')[0]

# 数据转换
data_transform = transforms.Compose([
    transforms.RandomResizedCrop(224, scale=(0.08, 1.0), ratio=(3.0/4.0, 4.0/3.0)),
    transforms.RandomHorizontalFlip(),
    transforms.ColorJitter(brightness=0.4, contrast=0.4, saturation=0.4, hue=0.2), # 增加 hue 抖动
    transforms.RandomRotation(20), # 增加旋转角度
    transforms.RandomAffine(degrees=0, translate=(0.1, 0.1), scale=(0.9, 1.1)), # 添加仿射变换
    transforms.GaussianBlur(kernel_size=(5, 9), sigma=(0.1, 2.0)),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])

# 验证集转换
val_transform = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])

# 加载数据集
full_dataset = DogBreedDataset(csv_file='/kaggle/input/dog-breed-identification/labels.csv',
                               root_dir='/kaggle/input/dog-breed-identification/train',
                               transform=data_transform)
                               # 定义ResNet152模型
def get_model(num_classes):
    model = models.resnet152(weights=models.ResNet152_Weights.DEFAULT)
    for param in model.parameters():
        param.requires_grad = True  # 冻结所有层

    # 解冻 layer1, layer2, layer3, 和 layer4
    for layer in [model.layer1, model.layer2, model.layer3, model.layer4]:
        for param in layer.parameters():
            param.requires_grad = True

    num_ftrs = model.fc.in_features
    model.fc = nn.Sequential(
        nn.Dropout(0.5),
        nn.Linear(num_ftrs, 2048),
        nn.ReLU(),
        nn.Dropout(0.5),
        nn.Linear(2048, num_classes)
    )
    return model

# def get_model(num_classes):
#     model = models.resnet34(weights=models.ResNet34_Weights.DEFAULT)
#     for param in model.parameters():
#         param.requires_grad = False
#     num_ftrs = model.fc.in_features
#     model.fc = nn.Linear(num_ftrs, num_classes)
#     return model

# 训练函数
def train_model(model, train_loader, val_loader, criterion, optimizer, scheduler, num_epochs, device, grad_clip=1.0):
    scaler = GradScaler()
    best_val_loss = float('inf')
    patience = 5
    counter = 0

    for epoch in range(num_epochs):
        model.train()
        train_loss = 0.0
        train_correct = 0
        train_total = 0
        
        pbar = tqdm(train_loader, desc=f"Epoch {epoch+1}/{num_epochs} - Training")
        for inputs, labels in pbar:
            inputs, labels = inputs.to(device), labels.to(device)
            
            optimizer.zero_grad()
            with autocast():
                outputs = model(inputs)
                loss = criterion(outputs, labels)
            
            if torch.isnan(loss):
                print(f"NaN loss detected in epoch {epoch+1}. Stopping training.")
                return best_val_loss

            scaler.scale(loss).backward()
            
            # Unscale the gradients before clipping
            scaler.unscale_(optimizer)
            
            # Gradient clipping
            torch.nn.utils.clip_grad_norm_(model.parameters(), grad_clip)
            
            # Gradient monitoring
            total_norm = 0
            for p in model.parameters():
                if p.grad is not None:
                    param_norm = p.grad.data.norm(2)
                    total_norm += param_norm.item() ** 2
            total_norm = total_norm ** 0.5
            
            scaler.step(optimizer)
            scaler.update()
            
            if scheduler is not None:
                scheduler.step()

            train_loss += loss.item() * inputs.size(0)
            _, predicted = outputs.max(1)
            train_total += labels.size(0)
            train_correct += predicted.eq(labels).sum().item()

            # Update progress bar
            pbar.set_postfix({
                'loss': f"{loss.item():.4f}",
                'acc': f"{train_correct/train_total:.4f}",
                'lr': f"{optimizer.param_groups[0]['lr']:.6f}",
                'grad_norm': f"{total_norm:.4f}"
            })

        train_loss /= len(train_loader.dataset)
        train_acc = train_correct / train_total
        
        # Validation
        val_loss, val_acc = validate_model(model, val_loader, criterion, device)
        
        print(f'Epoch {epoch+1}/{num_epochs}:')
        print(f'Train Loss: {train_loss:.4f} Acc: {train_acc:.4f}')
        print(f'Val Loss: {val_loss:.4f} Acc: {val_acc:.4f}')
        print(f'Learning Rate: {optimizer.param_groups[0]["lr"]:.6f}')
        print(f'Gradient Norm: {total_norm:.4f}')
        
        if val_loss < best_val_loss:
            best_val_loss = val_loss
            torch.save(model.state_dict(), 'best_model.pth')
            print("Saved best model.")
            counter = 0
        else:
            counter += 1
            if counter >= patience:
                print(f"Early stopping at epoch {epoch+1}")
                break
    
    return best_val_loss

def validate_model(model, val_loader, criterion, device):
    model.eval()
    val_loss = 0.0
    val_correct = 0
    val_total = 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)
            _, predicted = outputs.max(1)
            val_total += labels.size(0)
            val_correct += predicted.eq(labels).sum().item()
    
    val_loss /= len(val_loader.dataset)
    val_acc = val_correct / val_total
    return val_loss, val_acc

    # 主训练循环
def main():
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    num_classes = len(full_dataset.breeds)
    num_epochs = 200
    batch_size = 16
    train_size = int(0.8 * len(full_dataset))
    val_size = len(full_dataset) - train_size
    train_dataset, val_dataset = torch.utils.data.random_split(full_dataset, [train_size, val_size])
    
    # 为验证集应用不同的转换
    val_dataset.dataset.transform = val_transform
    
    train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=16)
    val_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False, num_workers=16)
    # 初始化模型、损失函数、优化器和调度器
    model = get_model(num_classes).to(device)
    criterion = nn.CrossEntropyLoss()

# 为不同层设置不同的学习率
    optimizer = optim.AdamW([
        {'params': model.fc.parameters(), 'lr': 5e-4},
        {'params': model.layer1.parameters(), 'lr': 1e-5},
        {'params': model.layer2.parameters(), 'lr': 2e-5},
        {'params': model.layer3.parameters(), 'lr': 3e-5},
        {'params': model.layer4.parameters(), 'lr': 5e-5}
    ], weight_decay=1e-2)

    scheduler = OneCycleLR(
        optimizer, 
        max_lr=[5e-4, 1e-5, 2e-5, 3e-5, 5e-5],
        steps_per_epoch=len(train_loader),
        epochs=num_epochs,
        pct_start=0.5,        # 增加预热阶段
        div_factor=100,       # 增加div_factor以降低初始学习率
        final_div_factor=500  # 减小final_div_factor使最终学习率不那么低
    )

    # 训练模型
    best_loss = train_model(model, train_loader, val_loader, criterion, optimizer, scheduler, num_epochs, device, grad_clip=1.0)

    print(f"Best Validation Loss: {best_loss:.4f}")
    # 加载最佳模型
    best_model = get_model(num_classes).to(device)
    best_model.load_state_dict(torch.load('best_model.pth'))
    return best_model

# 测试函数
def test(model, test_loader, device, class_names):
    predictions = []
    model.eval()

    with torch.no_grad():
        for inputs, image_ids in tqdm(test_loader):
            inputs = inputs.to(device)
            outputs = model(inputs)
            probs = torch.softmax(outputs, dim=1)
            probs = probs.cpu().numpy()  # 将概率转换为 NumPy 数组

            for image_id, prob in zip(image_ids, probs):
                prediction = {'id': image_id}
                for i, class_name in enumerate(class_names):
                    prediction[class_name] = prob[i]  # 将每个类别的概率添加到字典中
                predictions.append(prediction)

    return predictions

if __name__ == "__main__":
    best_model = main()
    
    # 准备测试数据
    test_dataset = TestDataset(root_dir='/kaggle/input/dog-breed-identification/test', transform=val_transform)
    test_loader = DataLoader(test_dataset, batch_size=32, shuffle=False, num_workers=16)
    
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    # 获取所有类别名称
    class_names = full_dataset.breeds.tolist()
    
    # 进行测试
    predictions = test(best_model, test_loader, device, class_names)

    # 创建 DataFrame 并保存为 CSV
    submission_df = pd.DataFrame(predictions, columns=['id'] + class_names)  # 设置列名
    submission_df.sort_values(by='id', inplace=True)
    submission_df.to_csv('/kaggle/working/submission.csv', index=False)

    print("Predictions saved to submission.csv")

此代码的loss较高，如何调整让其下降，注意不要用k折交叉验证，数据转换也足够复杂，不需要改了，OneCycleLR效果也比较好，不需要改了，从其他地方改进，给我完整的代码