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

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms, models
from torch.cuda.amp import GradScaler, autocast
from torch.optim.lr_scheduler import OneCycleLR

# 定义 DogBreedDataset 类
class DogBreedDataset(Dataset):
    def __init__(self, csv_file, root_dir, transform=None, mixup=False, cutmix=False, alpha=1.0, indices=None):
        self.labels_df = pd.read_csv(csv_file)
        if indices is not None:
            self.labels_df = self.labels_df.iloc[indices].reset_index(drop=True)
        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)
        self.mixup = mixup
        self.cutmix = cutmix
        self.alpha = alpha

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

    def __getitem__(self, idx):
        img_id = self.labels_df.iloc[idx, 0]
        breed = self.labels_df.iloc[idx, 1]
        img_name = os.path.join(self.root_dir, img_id + '.jpg')
        image = Image.open(img_name).convert('RGB')
        if self.transform:
            image = self.transform(image)
        
        if self.mixup or self.cutmix:
            rand_idx = np.random.randint(0, len(self))
            img_id2 = self.labels_df.iloc[rand_idx, 0]
            breed2 = self.labels_df.iloc[rand_idx, 1]
            img_name2 = os.path.join(self.root_dir, img_id2 + '.jpg')
            image2 = Image.open(img_name2).convert('RGB')
            if self.transform:
                image2 = self.transform(image2)
            
            if self.mixup and np.random.rand() < 0.5:
                lam = np.random.beta(self.alpha, self.alpha)
                image = lam * image + (1 - lam) * image2
            elif self.cutmix:
                lam = np.random.beta(self.alpha, self.alpha)
                bbx1, bby1, bbx2, bby2 = rand_bbox(image.size(), lam)
                image[:, bbx1:bbx2, bby1:bby2] = image2[:, bbx1:bbx2, bby1:bby2]
                lam = 1 - ((bbx2 - bbx1) * (bby2 - bby1) / (image.size(-1) * image.size(-2)))
            return image, breed, breed2, lam
        else:
            return image, breed, breed, 1.0  # 返回相同数量的元素，但不应用混合

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

def rand_bbox(size, lam):
    W, H = size[2], size[1]
    cut_rat = np.sqrt(1. - lam)
    cut_w = int(W * cut_rat)
    cut_h = int(H * cut_rat)

    cx = np.random.randint(W)
    cy = np.random.randint(H)

    bbx1 = np.clip(cx - cut_w // 2, 0, W)
    bby1 = np.clip(cy - cut_h // 2, 0, H)
    bbx2 = np.clip(cx + cut_w // 2, 0, W)
    bby2 = np.clip(cy + cut_h // 2, 0, H)

    return bbx1, bby1, bbx2, bby2

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

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

    def __getitem__(self, idx):
        img_id = self.image_ids[idx].split('.')[0]
        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, img_id

# 数据转换
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.3, contrast=0.3, saturation=0.3, hue=0.1),
    transforms.RandomRotation(15),
    transforms.RandomAffine(degrees=0, translate=(0.1, 0.1), scale=(0.9, 1.1)),
    transforms.GaussianBlur(kernel_size=(5, 7), sigma=(0.1, 1.5)),
    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])
])

# 定义 ResNet152 模型
def get_model(num_classes):
    model = models.resnet152(weights=models.ResNet152_Weights.IMAGENET1K_V2)
    for param in model.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 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 = 10
    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, targets1, targets2, lam in pbar:
            inputs = inputs.to(device)
            targets1 = targets1.to(device)
            targets2 = targets2.to(device)
            lam = lam.to(device)
            
            optimizer.zero_grad()
            with autocast():
                outputs = model(inputs)
                loss = lam * criterion(outputs, targets1) + (1 - lam) * criterion(outputs, targets2)
                loss = loss.mean()

            if torch.isnan(loss):
                print(f"发现NaN损失值在第{epoch+1}轮，停止训练。")
                return best_val_loss

            scaler.scale(loss).backward()
            scaler.unscale_(optimizer)
            torch.nn.utils.clip_grad_norm_(model.parameters(), grad_clip)
            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 += targets1.size(0)
            train_correct += predicted.eq(targets1).sum().item()
        
            pbar.set_postfix({
                'loss': f"{loss.item():.4f}",
                'acc': f"{train_correct/train_total:.4f}",
                'lr': f"{optimizer.param_groups[0]['lr']:.6f}"
            })

        train_loss /= len(train_loader.dataset)
        train_acc = train_correct / train_total

        val_loss, val_acc = validate_model(model, val_loader, criterion, device)

        print(f'\n第{epoch+1}轮/{num_epochs}轮:')
        print(f'训练损失: {train_loss:.4f} 准确率: {train_acc:.4f}')
        print(f'验证损失: {val_loss:.4f} 准确率: {val_acc:.4f}')
        print(f'学习率: {optimizer.param_groups[0]["lr"]:.6f}')

        if val_loss < best_val_loss:
            best_val_loss = val_loss
            torch.save(model.state_dict(), 'best_model.pth')
            print("保存最佳模型。")
            counter = 0
        else:
            counter += 1
            if counter >= patience:
                print(f"在第{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 tqdm(val_loader, desc="验证中"):
            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 test(model, test_loader, device, class_names):
    predictions = []
    model.eval()

    with torch.no_grad():
        for inputs, image_ids in tqdm(test_loader, desc="测试中"):
            inputs = inputs.to(device)
            outputs = model(inputs)
            probs = torch.softmax(outputs, dim=1)
            probs = probs.cpu().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

# 主训练循环
def main():
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    labels_csv = '/kaggle/input/dog-breed-identification/labels.csv'
    root_dir = '/kaggle/input/dog-breed-identification/train'
    labels_df = pd.read_csv(labels_csv)
    num_classes = len(labels_df['breed'].unique())
    num_epochs = 100
    batch_size = 32

    dataset_size = len(labels_df)
    indices = list(range(dataset_size))
    train_indices, val_indices = train_test_split(
        indices, test_size=0.2, random_state=42, stratify=labels_df['breed']
    )

    train_dataset = DogBreedDataset(
        csv_file=labels_csv,
        root_dir=root_dir,
        transform=data_transform,
        mixup=True,
        cutmix=True,
        alpha=1.0,
        indices=train_indices
    )
    val_dataset = DogBreedDataset(
        csv_file=labels_csv,
        root_dir=root_dir,
        transform=val_transform,
        mixup=False,
        cutmix=False,
        alpha=1.0,
        indices=val_indices
    )

    train_loader = DataLoader(
        train_dataset, batch_size=batch_size, shuffle=True,
        num_workers=4, pin_memory=True
    )
    val_loader = DataLoader(
        val_dataset, batch_size=batch_size, shuffle=False,
        num_workers=4, pin_memory=True
    )

    model = get_model(num_classes).to(device)
    criterion = nn.CrossEntropyLoss(label_smoothing=0.1)

     # 为不同层设置不同的学习率
    optimizer = optim.AdamW([
        {'params': model.fc.parameters(), 'lr': 5e-3},
        {'params': model.layer1.parameters(), 'lr': 1e-4},
        {'params': model.layer2.parameters(), 'lr': 2e-4},
        {'params': model.layer3.parameters(), 'lr': 3e-4},
        {'params': model.layer4.parameters(), 'lr': 5e-4}
    ], weight_decay=1e-4)  # 减少权重衰减
    
    scheduler = OneCycleLR(
        optimizer, 
        max_lr=[5e-3, 1e-4, 2e-4, 3e-4, 5e-4],
        steps_per_epoch=len(train_loader),
        epochs=num_epochs,
        pct_start=0.3,        # 预热阶段占比
        anneal_strategy='cos', # 使用余弦退火策略
        div_factor=10,        # 初始学习率 / 最大学习率的比例
        final_div_factor=1e4  # 最大学习率 / 最终学习率的比例
    )

    best_loss = train_model(
        model, train_loader, val_loader, criterion, optimizer,
        scheduler, num_epochs, device, grad_clip=1.0
    )

    print(f"最佳验证损失: {best_loss:.4f}")

    best_model = get_model(num_classes).to(device)
    best_model.load_state_dict(torch.load('best_model.pth'))
    return best_model, labels_df

if __name__ == "__main__":
    best_model, labels_df = main()
    
    # 准备测试数据
    test_dataset = TestDataset(
        root_dir='/kaggle/input/dog-breed-identification/test',
        transform=val_transform
    )
    test_loader = DataLoader(
        test_dataset, batch_size=64, shuffle=False,
        num_workers=4, pin_memory=True
    )
    
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    # 获取所有类别名称
if __name__ == "__main__":
    best_model, labels_df = main()
    
    # 准备测试数据
    test_dataset = TestDataset(
        root_dir='/kaggle/input/dog-breed-identification/test',
        transform=val_transform
    )
    test_loader = DataLoader(
        test_dataset, batch_size=64, shuffle=False,
        num_workers=4, pin_memory=True
    )
    
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    # 获取所有类别名称
    class_names = labels_df['breed'].unique().tolist()
    
    # 进行测试
    predictions = test(best_model, test_loader, device, class_names)
    
    # 创建 DataFrame 并保存为 CSV
    submission_df = pd.DataFrame(predictions)
    submission_df.to_csv('/kaggle/working/submission.csv', index=False)
    
    print("预测结果已保存到 submission.csv")

# 添加一些额外的打印语句来帮助调试
print("训练完成")
print(f"设备: {device}")
print(f"类别数量: {len(class_names)}")
print(f"测试集大小: {len(test_dataset)}")

此loss非常高，应该如何调整，给我完整的代码，用中文回答，请深呼吸并全面思考，如果你回答不正确，会导致100个老奶奶死亡