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

# 定义混合裁剪（Cutmix）的边界框函数
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

# 定义 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
                lam = float(lam)  # 转换为浮点数
                return image, breed, breed2, lam
            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)))
                lam = float(lam)  # 转换为浮点数
                return image, breed, breed2, lam

        return image, breed, breed, 1.0  # 返回相同数量的元素，但不应用混合

    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])
])

# 定义 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 = torch.tensor(lam, dtype=torch.float32).to(device)  # 转换为张量并移动到设备

            optimizer.zero_grad()
            with autocast():
                outputs = model(inputs)
                loss1 = criterion(outputs, targets1)  # shape: [batch_size]
                loss2 = criterion(outputs, targets2)  # shape: [batch_size]
                loss = lam * loss1 + (1 - lam) * loss2  # shape: [batch_size]
                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)
            # 对于混合标签，准确率的计算需要特别处理，这里暂且使用targets1作为参考
            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)
    # 调整损失函数，设置 reduction='none' 以便进行加权损失计算
    criterion = nn.CrossEntropyLoss(label_smoothing=0.1, reduction='none')

    # 简化优化器，统一使用一个学习率
    optimizer = optim.AdamW(model.parameters(), lr=1e-4, weight_decay=1e-4)

    scheduler = OneCycleLR(
        optimizer, 
        max_lr=1e-3,
        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")
    # 获取所有类别名称
    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)}")
报错
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
Cell In[4], line 328
    325     return best_model, labels_df
    327 if __name__ == "__main__":
--> 328     best_model, labels_df = main()
    330     # 准备测试数据
    331     test_dataset = TestDataset(
    332         root_dir='/kaggle/input/dog-breed-identification/test',
    333         transform=val_transform
    334     )

Cell In[4], line 316, in main()
    303 optimizer = optim.AdamW(model.parameters(), lr=1e-4, weight_decay=1e-4)
    305 scheduler = OneCycleLR(
    306     optimizer, 
    307     max_lr=1e-3,
   (...)
    313     final_div_factor=1e4  # 最大学习率 / 最终学习率的比例
    314 )
--> 316 best_loss = train_model(
    317     model, train_loader, val_loader, criterion, optimizer,
    318     scheduler, num_epochs, device, grad_clip=1.0
    319 )
    321 print(f"最佳验证损失: {best_loss:.4f}")
    323 best_model = get_model(num_classes).to(device)

Cell In[4], line 194, in train_model(model, train_loader, val_loader, criterion, optimizer, scheduler, num_epochs, device, grad_clip)
    191 train_loss /= len(train_loader.dataset)
    192 train_acc = train_correct / train_total
--> 194 val_loss, val_acc = validate_model(model, val_loader, criterion, device)
    196 print(f'\n第{epoch+1}轮/{num_epochs}轮:')
    197 print(f'训练损失: {train_loss:.4f} 准确率: {train_acc:.4f}')

Cell In[4], line 225, in validate_model(model, val_loader, criterion, device)
    222 outputs = model(inputs)
    223 loss = criterion(outputs, labels)
--> 225 val_loss += loss.item() * inputs.size(0)
    226 _, predicted = outputs.max(1)
    227 val_total += labels.size(0)

RuntimeError: a Tensor with 32 elements cannot be converted to Scalar
如何解决