def get_model(num_classes):
    model = models.resnet152(weights=models.ResNet152_Weights.DEFAULT)
    for param in model.parameters():
        param.requires_grad = True

    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 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, labels2, lam in pbar:
            inputs, labels, labels2, lam = inputs.to(device), labels.to(device), labels2.to(device), lam.to(device)
            
            optimizer.zero_grad()
            with autocast():
                outputs = model(inputs)
                loss = lam * criterion(outputs, labels) + (1 - lam) * criterion(outputs, labels2)
                loss = loss.mean()  # 确保loss是一个标量
            
            if torch.isnan(loss):
                print(f"NaN loss detected in epoch {epoch+1}. Stopping training.")
                return best_val_loss

            scaler.scale(loss).backward()
            
            scaler.unscale_(optimizer)
            torch.nn.utils.clip_grad_norm_(model.parameters(), grad_clip)
            
            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()

            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
        
        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")
    labels_csv = '/kaggle/input/dog-breed-identification/labels.csv'
    root_dir = '/kaggle/input/dog-breed-identification/train'
    full_dataset = DogBreedDataset(csv_file=labels_csv, root_dir=root_dir, transform=data_transform)
    num_classes = len(full_dataset.breeds)
    class_names = full_dataset.breeds.tolist()
    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(label_smoothing=0.1)

    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)

    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,
        final_div_factor=500
    )

    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, class_names, num_classes

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

            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, class_names, num_classes = 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")
    
    predictions = test(best_model, test_loader, device, class_names)

    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")
    print("Training completed")
    print(f"Device: {device}")
    print(f"Number of classes: {num_classes}")
    print(f"Test set size: {len(test_dataset)}")
loss太高，如何通过调整参数的方式降低，不要修改代码的实现