pip install torch torchvision torchaudio matplotlib scikit-learn pandas tqdm

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, random_split
from torchvision import datasets, transforms, models
from torchvision.models import ResNet50_Weights
import matplotlib.pyplot as plt
from sklearn.metrics import classification_report, confusion_matrix
import numpy as np
import os
from tqdm import tqdm

# ==================== 配置參數(shù) ====================
class Config:
    num_classes = 10
    batch_size = 64
    num_epochs = 20
    learning_rate = 1e-3
    weight_decay = 1e-4
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    save_path = 'best_model.pth'
    num_workers = 4

config = Config()

# ==================== 數(shù)據(jù)預(yù)處理 ====================
def get_transforms():
    """獲取訓(xùn)練和驗證的變換"""
    train_transform = transforms.Compose([
        transforms.RandomResizedCrop(224, scale=(0.8, 1.0)),
        transforms.RandomHorizontalFlip(p=0.5),
        transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2, hue=0.1),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
    ])
    
    val_transform = transforms.Compose([
        transforms.Resize(256),
        transforms.CenterCrop(224),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
    ])
    
    return train_transform, val_transform

# ==================== 數(shù)據(jù)加載 ====================
def load_data():
    """加載并分割數(shù)據(jù)集"""
    train_transform, val_transform = get_transforms()
    
    # 加載 CIFAR-10 數(shù)據(jù)集
    full_train_dataset = datasets.CIFAR10(
        root='./data', train=True, download=True, transform=train_transform
    )
    test_dataset = datasets.CIFAR10(
        root='./data', train=False, download=True, transform=val_transform
    )
    
    # 分割訓(xùn)練集為訓(xùn)練集和驗證集 (90:10)
    train_size = int(0.9 * len(full_train_dataset))
    val_size = len(full_train_dataset) - train_size
    train_dataset, val_dataset = random_split(
        full_train_dataset, [train_size, val_size]
    )
    
    # 創(chuàng)建數(shù)據(jù)加載器
    train_loader = DataLoader(
        train_dataset, batch_size=config.batch_size, shuffle=True,
        num_workers=config.num_workers, pin_memory=True
    )
    val_loader = DataLoader(
        val_dataset, batch_size=config.batch_size, shuffle=False,
        num_workers=config.num_workers, pin_memory=True
    )
    test_loader = DataLoader(
        test_dataset, batch_size=config.batch_size, shuffle=False,
        num_workers=config.num_workers, pin_memory=True
    )
    
    return train_loader, val_loader, test_loader

# ==================== 模型定義 ====================
class CustomResNet50(nn.Module):
    def __init__(self, num_classes=10, pretrained=True):
        super().__init__()
        if pretrained:
            weights = ResNet50_Weights.IMAGENET1K_V2
            self.model = models.resnet50(weights=weights)
        else:
            self.model = models.resnet50(weights=None)
        
        # 替換最后的全連接層
        num_features = self.model.fc.in_features
        self.model.fc = nn.Sequential(
            nn.Dropout(0.5),
            nn.Linear(num_features, num_classes)
        )
        
    def forward(self, x):
        return self.model(x)

# ==================== 訓(xùn)練函數(shù) ====================
def train_one_epoch(model, dataloader, criterion, optimizer, scheduler, device):
    """訓(xùn)練一個 epoch"""
    model.train()
    running_loss = 0.0
    correct = 0
    total = 0
    
    progress_bar = tqdm(dataloader, desc="Training")
    for inputs, targets in progress_bar:
        inputs, targets = inputs.to(device), targets.to(device)
        
        optimizer.zero_grad()
        outputs = model(inputs)
        loss = criterion(outputs, targets)
        loss.backward()
        
        # 梯度裁剪
        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
        
        optimizer.step()
        
        running_loss += loss.item()
        _, predicted = outputs.max(1)
        total += targets.size(0)
        correct += predicted.eq(targets).sum().item()
        
        progress_bar.set_postfix({
            'loss': loss.item(),
            'acc': 100. * correct / total
        })
    
    epoch_loss = running_loss / len(dataloader)
    epoch_acc = 100. * correct / total
    return epoch_loss, epoch_acc

# ==================== 驗證函數(shù) ====================
def validate(model, dataloader, criterion, device):
    """驗證模型"""
    model.eval()
    running_loss = 0.0
    correct = 0
    total = 0
    
    with torch.no_grad():
        for inputs, targets in dataloader:
            inputs, targets = inputs.to(device), targets.to(device)
            outputs = model(inputs)
            loss = criterion(outputs, targets)
            
            running_loss += loss.item()
            _, predicted = outputs.max(1)
            total += targets.size(0)
            correct += predicted.eq(targets).sum().item()
    
    epoch_loss = running_loss / len(dataloader)
    epoch_acc = 100. * correct / total
    return epoch_loss, epoch_acc

# ==================== 訓(xùn)練主循環(huán) ====================
def train_model():
    """完整的訓(xùn)練流程"""
    # 設(shè)置隨機種子
    torch.manual_seed(42)
    np.random.seed(42)
    
    # 加載數(shù)據(jù)
    print("Loading data...")
    train_loader, val_loader, test_loader = load_data()
    print(f"Train samples: {len(train_loader.dataset)}")
    print(f"Val samples: {len(val_loader.dataset)}")
    print(f"Test samples: {len(test_loader.dataset)}")
    
    # 初始化模型
    print("Initializing model...")
    model = CustomResNet50(num_classes=config.num_classes, pretrained=True)
    model = model.to(config.device)
    
    # 損失函數(shù)和優(yōu)化器
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.AdamW(
        model.parameters(), 
        lr=config.learning_rate, 
        weight_decay=config.weight_decay
    )
    
    # 學(xué)習(xí)率調(diào)度器
    scheduler = optim.lr_scheduler.CosineAnnealingLR(
        optimizer, T_max=config.num_epochs
    )
    
    # 訓(xùn)練歷史記錄
    train_losses, train_accs = [], []
    val_losses, val_accs = [], []
    best_val_acc = 0.0
    
    # 訓(xùn)練循環(huán)
    print(f"Starting training on {config.device}...")
    for epoch in range(config.num_epochs):
        print(f"\nEpoch {epoch+1}/{config.num_epochs}")
        
        # 訓(xùn)練
        train_loss, train_acc = train_one_epoch(
            model, train_loader, criterion, optimizer, scheduler, config.device
        )
        
        # 驗證
        val_loss, val_acc = validate(model, val_loader, criterion, config.device)
        
        # 更新學(xué)習(xí)率
        scheduler.step()
        
        # 記錄歷史
        train_losses.append(train_loss)
        train_accs.append(train_acc)
        val_losses.append(val_loss)
        val_accs.append(val_acc)
        
        # 保存最佳模型
        if val_acc > best_val_acc:
            best_val_acc = val_acc
            torch.save(model.state_dict(), config.save_path)
            print(f"Saved best model with validation accuracy: {best_val_acc:.2f}%")
        
        print(f"Train Loss: {train_loss:.4f}, Train Acc: {train_acc:.2f}%")
        print(f"Val Loss: {val_loss:.4f}, Val Acc: {val_acc:.2f}%")
    
    # 繪制訓(xùn)練曲線
    plot_training_history(train_losses, val_losses, train_accs, val_accs)
    
    # 測試最佳模型
    test_model(test_loader)
    
    return model

# ==================== 可視化函數(shù) ====================
def plot_training_history(train_losses, val_losses, train_accs, val_accs):
    """繪制訓(xùn)練歷史"""
    fig, axes = plt.subplots(1, 2, figsize=(12, 4))
    
    # 損失曲線
    axes[0].plot(train_losses, label='Train Loss')
    axes[0].plot(val_losses, label='Val Loss')
    axes[0].set_title('Training and Validation Loss')
    axes[0].set_xlabel('Epoch')
    axes[0].set_ylabel('Loss')
    axes[0].legend()
    axes[0].grid(True)
    
    # 準確率曲線
    axes[1].plot(train_accs, label='Train Accuracy')
    axes[1].plot(val_accs, label='Val Accuracy')
    axes[1].set_title('Training and Validation Accuracy')
    axes[1].set_xlabel('Epoch')
    axes[1].set_ylabel('Accuracy (%)')
    axes[1].legend()
    axes[1].grid(True)
    
    plt.tight_layout()
    plt.savefig('training_history.png')
    plt.show()

# ==================== 測試函數(shù) ====================
def test_model(test_loader):
    """測試模型性能"""
    model = CustomResNet50(num_classes=config.num_classes, pretrained=False)
    model.load_state_dict(torch.load(config.save_path))
    model = model.to(config.device)
    model.eval()
    
    all_preds = []
    all_targets = []
    
    with torch.no_grad():
        for inputs, targets in test_loader:
            inputs, targets = inputs.to(config.device), targets.to(config.device)
            outputs = model(inputs)
            _, preds = outputs.max(1)
            
            all_preds.extend(preds.cpu().numpy())
            all_targets.extend(targets.cpu().numpy())
    
    # 計算指標(biāo)
    accuracy = 100. * sum(np.array(all_preds) == np.array(all_targets)) / len(all_targets)
    print(f"\nTest Accuracy: {accuracy:.2f}%")
    
    # 分類報告
    class_names = ['airplane', 'automobile', 'bird', 'cat', 'deer', 
                   'dog', 'frog', 'horse', 'ship', 'truck']
    print("\nClassification Report:")
    print(classification_report(all_targets, all_preds, target_names=class_names))
    
    # 混淆矩陣
    cm = confusion_matrix(all_targets, all_preds)
    plt.figure(figsize=(10, 8))
    plt.imshow(cm, interpolation='nearest', cmap=plt.cm.Blues)
    plt.title('Confusion Matrix')
    plt.colorbar()
    tick_marks = np.arange(len(class_names))
    plt.xticks(tick_marks, class_names, rotation=45)
    plt.yticks(tick_marks, class_names)
    plt.tight_layout()
    plt.savefig('confusion_matrix.png')
    plt.show()

# ==================== 推理函數(shù) ====================
def predict_image(image_path, model_path='best_model.pth'):
    """對單張圖像進行預(yù)測"""
    # 加載模型
    model = CustomResNet50(num_classes=config.num_classes, pretrained=False)
    model.load_state_dict(torch.load(model_path))
    model = model.to(config.device)
    model.eval()
    
    # 圖像預(yù)處理
    transform = transforms.Compose([
        transforms.Resize(256),
        transforms.CenterCrop(224),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
    ])
    
    from PIL import Image
    image = Image.open(image_path).convert('RGB')
    input_tensor = transform(image).unsqueeze(0).to(config.device)
    
    # 預(yù)測
    with torch.no_grad():
        output = model(input_tensor)
        probabilities = torch.softmax(output, dim=1)
        confidence, predicted_class = torch.max(probabilities, dim=1)
    
    class_names = ['airplane', 'automobile', 'bird', 'cat', 'deer', 
                   'dog', 'frog', 'horse', 'ship', 'truck']
    
    result = {
        'predicted_class': class_names[predicted_class.item()],
        'confidence': confidence.item(),
        'all_probabilities': {class_names[i]: prob.item() 
                             for i, prob in enumerate(probabilities[0])}
    }
    
    return result

if __name__ == "__main__":
    # 訓(xùn)練模型
    trained_model = train_model()
    
    # 示例：預(yù)測單張圖像（需要替換為實際圖像路徑）
    # result = predict_image('path/to/your/image.jpg')
    # print(f"Predicted: {result['predicted_class']} (Confidence: {result['confidence']:.2f})")

transforms.RandomResizedCrop(224, scale=(0.8, 1.0))

transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2, hue=0.1)

torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)

# 余弦退火調(diào)度器
scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=num_epochs)

from torch.cuda.amp import autocast, GradScaler

scaler = GradScaler()

for inputs, targets in dataloader:
    optimizer.zero_grad()
    
    with autocast():
        outputs = model(inputs)
        loss = criterion(outputs, targets)
    
    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()

# 單機多卡
model = nn.DataParallel(model)

# 多機多卡 (DDP)
import torch.distributed as dist
dist.init_process_group(backend='nccl')
model = nn.parallel.DistributedDataParallel(model)

# 動態(tài)量化
quantized_model = torch.quantization.quantize_dynamic(
    model, {nn.Linear, nn.Conv2d}, dtype=torch.qint8
)

# 靜態(tài)量化
model.eval()
model.qconfig = torch.quantization.get_default_qconfig('fbgemm')
torch.quantization.prepare(model, inplace=True)
# 校準步驟...
torch.quantization.convert(model, inplace=True)

from transformers import ViTForImageClassification, ViTImageProcessor
from transformers import TrainingArguments, Trainer

# 加載預(yù)訓(xùn)練 ViT
model = ViTForImageClassification.from_pretrained(
    'google/vit-base-patch16-224',
    num_labels=10,
    ignore_mismatched_sizes=True
)

processor = ViTImageProcessor.from_pretrained('google/vit-base-patch16-224')

# 自定義數(shù)據(jù)集
class CIFAR10Dataset(torch.utils.data.Dataset):
    def __init__(self, dataset, processor):
        self.dataset = dataset
        self.processor = processor
    
    def __len__(self):
        return len(self.dataset)
    
    def __getitem__(self, idx):
        image, label = self.dataset[idx]
        encoding = self.processor(image, return_tensors='pt')
        return {
            'pixel_values': encoding['pixel_values'].squeeze(),
            'labels': label
        }

# 訓(xùn)練配置
training_args = TrainingArguments(
    output_dir='./vit_results',
    per_device_train_batch_size=32,
    per_device_eval_batch_size=32,
    evaluation_strategy="epoch",
    num_train_epochs=5,
    fp16=True,
    save_steps=100,
    save_total_limit=2,
    logging_steps=10,
)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=train_dataset,
    eval_dataset=val_dataset,
    tokenizer=processor,
)

trainer.train()