import os
import pandas as pd
from PIL import Image
from tqdm import tqdm

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):
        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_id = self.labels_df.iloc[idx, 0]
        img_name = os.path.join(self.root_dir, img_id + '.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 = 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])
])

# 加载数据集
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
    
    # 替换全连接层
    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 = 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损失值在第{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 += labels.size(0)
            train_correct += predicted.eq(labels).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 main():
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    num_classes = len(full_dataset.breeds)
    num_epochs = 100  # 根据需要调整
    batch_size = 32     # 根据显存情况调整
    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=8, pin_memory=True)
    val_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False, num_workers=8, pin_memory=True)
    
    # 初始化模型、损失函数、优化器和调度器
    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-4)  # 减少权重衰减
    
    # 计算总训练步数
    total_steps = len(train_loader) * num_epochs
    
    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.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

# 测试函数
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()  # 将概率转换为 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=64, shuffle=False, num_workers=8, pin_memory=True)
    
    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("预测结果已保存到 submission.csv")
如何调整，才能降低loss，给我完整的代码，用中文回答