PyTorch：现代深度学习的核心框架

import torch
import torch.nn as nn

# 动态计算图的灵活性体现
def dynamic_model(x, use_dropout=True):
    x = nn.Linear(128, 64)(x)
    if use_dropout:  # 根据运行时条件动态构建图
        x = nn.Dropout(0.5)(x)
    x = torch.relu(x)
    return x

import torch

# 创建各种类型的张量
scalar = torch.tensor(3.14)           # 标量（0维张量）
vector = torch.tensor([1, 2, 3])      # 向量（1维张量）
matrix = torch.ones(3, 4)             # 矩阵（2维张量）
tensor_3d = torch.zeros(2, 3, 4)      # 3维张量

# GPU加速
if torch.cuda.is_available():
    device = torch.device('cuda')
    matrix_gpu = matrix.to(device)
    print(f"张量位于: {matrix_gpu.device}")

# 张量操作
a = torch.randn(3, 4)
b = torch.randn(4, 5)
c = torch.matmul(a, b)  # 矩阵乘法
print(f"结果形状: {c.shape}")  # torch.Size([3, 5])

import torch

# 创建需要梯度的张量
x = torch.tensor(2.0, requires_grad=True)
y = torch.tensor(3.0, requires_grad=True)

# 定义计算过程
z = x**2 + 2*x*y + y**3

# 反向传播计算梯度
z.backward()

print(f"dz/dx = {x.grad}")  # 2x + 2y = 4 + 6 = 10
print(f"dz/dy = {y.grad}")  # 2x + 3y² = 4 + 27 = 31

import torch
import torch.nn as nn
import torch.nn.functional as F

class ConvNet(nn.Module):
    def __init__(self, num_classes=10):
        super(ConvNet, self).__init__()
        # 定义网络层
        self.conv1 = nn.Conv2d(1, 32, kernel_size=3, padding=1)
        self.conv2 = nn.Conv2d(32, 64, kernel_size=3, padding=1)
        self.pool = nn.MaxPool2d(2, 2)
        self.dropout = nn.Dropout(0.25)
        self.fc1 = nn.Linear(64 * 7 * 7, 128)
        self.fc2 = nn.Linear(128, num_classes)
        self.bn1 = nn.BatchNorm2d(32)
        self.bn2 = nn.BatchNorm2d(64)
    
    def forward(self, x):
        # 前向传播逻辑
        x = self.pool(F.relu(self.bn1(self.conv1(x))))
        x = self.pool(F.relu(self.bn2(self.conv2(x))))
        x = self.dropout(x)
        x = x.view(-1, 64 * 7 * 7)  # 展平
        x = F.relu(self.fc1(x))
        x = self.dropout(x)
        x = self.fc2(x)
        return x

# 实例化模型
model = ConvNet(num_classes=10)
print(model)

# 统计参数量
total_params = sum(p.numel() for p in model.parameters())
trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
print(f"总参数量: {total_params:,}")
print(f"可训练参数量: {trainable_params:,}")

import torch
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms, datasets
import numpy as np

# 自定义数据集
class CustomDataset(Dataset):
    def __init__(self, data_path, transform=None):
        self.data = np.load(data_path)  # 加载数据
        self.transform = transform
    
    def __len__(self):
        return len(self.data)
    
    def __getitem__(self, idx):
        sample = self.data[idx]
        if self.transform:
            sample = self.transform(sample)
        return sample

# 使用torchvision内置数据集（以CIFAR-10为例）
transform_train = transforms.Compose([
    transforms.RandomCrop(32, padding=4),
    transforms.RandomHorizontalFlip(),
    transforms.ColorJitter(brightness=0.2, contrast=0.2),
    transforms.ToTensor(),
    transforms.Normalize((0.4914, 0.4822, 0.4465),
                         (0.2023, 0.1994, 0.2010))
])

transform_val = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.4914, 0.4822, 0.4465),
                         (0.2023, 0.1994, 0.2010))
])

train_dataset = datasets.CIFAR10(
    root='./data', train=True, 
    download=True, transform=transform_train
)
val_dataset = datasets.CIFAR10(
    root='./data', train=False, 
    download=True, transform=transform_val
)

# 创建DataLoader
train_loader = DataLoader(
    train_dataset,
    batch_size=128,
    shuffle=True,
    num_workers=4,      # 多进程数据加载
    pin_memory=True,    # 锁页内存，加速GPU传输
    drop_last=True      # 丢弃不完整的最后一批
)

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

print(f"训练集大小: {len(train_dataset)}")
print(f"验证集大小: {len(val_dataset)}")
print(f"训练批次数: {len(train_loader)}")

import torch
import torch.nn as nn
import torch.optim as optim
from torch.optim.lr_scheduler import CosineAnnealingLR
import time

def train_epoch(model, loader, criterion, optimizer, device, epoch):
    """单个训练轮次"""
    model.train()
    total_loss = 0.0
    correct = 0
    total = 0
    
    for batch_idx, (inputs, targets) in enumerate(loader):
        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()
        
        # 统计指标
        total_loss += loss.item()
        _, predicted = outputs.max(1)
        total += targets.size(0)
        correct += predicted.eq(targets).sum().item()
        
        if batch_idx % 100 == 0:
            print(f'Epoch [{epoch}] Batch [{batch_idx}/{len(loader)}] '
                  f'Loss: {loss.item():.4f} '
                  f'Acc: {100.*correct/total:.2f}%')
    
    return total_loss / len(loader), 100. * correct / total


@torch.no_grad()
def evaluate(model, loader, criterion, device):
    """模型评估"""
    model.eval()
    total_loss = 0.0
    correct = 0
    total = 0
    
    for inputs, targets in loader:
        inputs, targets = inputs.to(device), targets.to(device)
        outputs = model(inputs)
        loss = criterion(outputs, targets)
        
        total_loss += loss.item()
        _, predicted = outputs.max(1)
        total += targets.size(0)
        correct += predicted.eq(targets).sum().item()
    
    return total_loss / len(loader), 100. * correct / total


def train_model(model, train_loader, val_loader, num_epochs=100):
    """完整训练流程"""
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    model = model.to(device)
    
    # 定义损失函数和优化器
    criterion = nn.CrossEntropyLoss(label_smoothing=0.1)
    optimizer = optim.AdamW(
        model.parameters(),
        lr=0.001,
        weight_decay=0.01,
        betas=(0.9, 0.999)
    )
    
    # 学习率调度
    scheduler = CosineAnnealingLR(optimizer, T_max=num_epochs, eta_min=1e-6)
    
    best_val_acc = 0.0
    history = {'train_loss': [], 'train_acc': [], 'val_loss': [], 'val_acc': []}
    
    for epoch in range(1, num_epochs + 1):
        start_time = time.time()
        
        train_loss, train_acc = train_epoch(
            model, train_loader, criterion, optimizer, device, epoch
        )
        val_loss, val_acc = evaluate(model, val_loader, criterion, device)
        scheduler.step()
        
        elapsed = time.time() - start_time
        
        # 记录历史
        history['train_loss'].append(train_loss)
        history['train_acc'].append(train_acc)
        history['val_loss'].append(val_loss)
        history['val_acc'].append(val_acc)
        
        print(f'\nEpoch [{epoch}/{num_epochs}] '
              f'Train Loss: {train_loss:.4f} Train Acc: {train_acc:.2f}% '
              f'Val Loss: {val_loss:.4f} Val Acc: {val_acc:.2f}% '
              f'LR: {scheduler.get_last_lr()[0]:.6f} '
              f'Time: {elapsed:.1f}s')
        
        # 保存最佳模型
        if val_acc > best_val_acc:
            best_val_acc = val_acc
            torch.save({
                'epoch': epoch,
                'model_state_dict': model.state_dict(),
                'optimizer_state_dict': optimizer.state_dict(),
                'val_acc': val_acc,
            }, 'best_model.pth')
            print(f'✓ 保存最佳模型 (Val Acc: {val_acc:.2f}%)')
    
    print(f'\n训练完成！最佳验证准确率: {best_val_acc:.2f}%')
    return history

import torch

# 保存完整模型（推荐保存state_dict）
def save_checkpoint(model, optimizer, epoch, filepath):
    checkpoint = {
        'epoch': epoch,
        'model_state_dict': model.state_dict(),
        'optimizer_state_dict': optimizer.state_dict(),
        'model_config': model.config if hasattr(model, 'config') else None
    }
    torch.save(checkpoint, filepath)
    print(f"检查点已保存至: {filepath}")

# 加载模型
def load_checkpoint(filepath, model, optimizer=None):
    checkpoint = torch.load(filepath, map_location='cpu')
    model.load_state_dict(checkpoint['model_state_dict'])
    
    if optimizer and 'optimizer_state_dict' in checkpoint:
        optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
    
    epoch = checkpoint.get('epoch', 0)
    print(f"从第 {epoch} 轮恢复训练")
    return model, optimizer, epoch

# 推理时使用torch.no_grad()
def inference(model, input_data, device):
    model.eval()
    with torch.no_grad():
        input_tensor = torch.tensor(input_data).to(device)
        output = model(input_tensor)
        probabilities = torch.softmax(output, dim=-1)
        predicted_class = torch.argmax(probabilities, dim=-1)
    return predicted_class.cpu().numpy()

import torch
from torch.cuda.amp import autocast, GradScaler

def train_with_amp(model, loader, criterion, optimizer, device):
    """使用自动混合精度训练"""
    model.train()
    scaler = GradScaler()  # 梯度缩放器，防止FP16下溢
    
    for inputs, targets in loader:
        inputs, targets = inputs.to(device), targets.to(device)
        optimizer.zero_grad()
        
        # 在autocast上下文中执行前向传播
        with autocast():
            outputs = model(inputs)
            loss = criterion(outputs, targets)
        
        # 使用scaler进行反向传播
        scaler.scale(loss).backward()
        scaler.unscale_(optimizer)
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        scaler.step(optimizer)
        scaler.update()
    
    return loss.item()

import torch
import torch.distributed as dist
import torch.nn as nn
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data.distributed import DistributedSampler
import os

def setup_distributed(rank, world_size):
    """初始化分布式训练环境"""
    os.environ['MASTER_ADDR'] = 'localhost'
    os.environ['MASTER_PORT'] = '12355'
    dist.init_process_group(
        backend='nccl',  # GPU通信使用NCCL
        rank=rank,
        world_size=world_size
    )
    torch.cuda.set_device(rank)

def cleanup_distributed():
    dist.destroy_process_group()

def train_distributed(rank, world_size, model, dataset):
    """分布式训练主函数"""
    setup_distributed(rank, world_size)
    
    # 将模型包装为DDP
    model = model.to(rank)
    ddp_model = DDP(model, device_ids=[rank])
    
    # 使用DistributedSampler确保数据均匀分布
    sampler = DistributedSampler(
        dataset, 
        num_replicas=world_size, 
        rank=rank
    )
    loader = DataLoader(dataset, batch_size=64, sampler=sampler)
    
    criterion = nn.CrossEntropyLoss()
    optimizer = torch.optim.Adam(ddp_model.parameters(), lr=0.001)
    
    for epoch in range(100):
        sampler.set_epoch(epoch)  # 确保每轮数据shuffle不同
        for inputs, targets in loader:
            inputs = inputs.to(rank)
            targets = targets.to(rank)
            
            optimizer.zero_grad()
            outputs = ddp_model(inputs)
            loss = criterion(outputs, targets)
            loss.backward()
            optimizer.step()
    
    cleanup_distributed()

# 启动多进程训练
import torch.multiprocessing as mp

if __name__ == '__main__':
    world_size = torch.cuda.device_count()
    mp.spawn(
        train_distributed,
        args=(world_size, model, dataset),
        nprocs=world_size,
        join=True
    )

import torch

# 使用torch.compile加速模型
model = MyModel()

# 不同的编译模式
compiled_model = torch.compile(model)  # 默认模式

# 最大性能优化（编译时间较长）
compiled_model_max = torch.compile(model, mode='max-autotune')

# 减少编译开销
compiled_model_reduce = torch.compile(model, mode='reduce-overhead')

# 测试性能提升
import time

x = torch.randn(32, 3, 224, 224).cuda()

# 原始模型
model.cuda().eval()
with torch.no_grad():
    start = time.time()
    for _ in range(100):
        _ = model(x)
    torch.cuda.synchronize()
    original_time = time.time() - start

# 编译后模型
compiled_model.cuda().eval()
with torch.no_grad():
    # 预热
    for _ in range(10):
        _ = compiled_model(x)
    
    start = time.time()
    for _ in range(100):
        _ = compiled_model(x)
    torch.cuda.synchronize()
    compiled_time = time.time() - start

print(f"原始模型时间: {original_time:.3f}s")
print(f"编译后模型时间: {compiled_time:.3f}s")
print(f"性能提升: {original_time/compiled_time:.2f}x")

import torch
import torch.nn as nn

class SimpleModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.linear = nn.Linear(10, 5)
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return torch.relu(self.linear(x))

model = SimpleModel()
model.eval()

# 方法一：Script（适合含控制流的模型）
scripted_model = torch.jit.script(model)
scripted_model.save('model_scripted.pt')

# 方法二：Trace（适合固定计算图的模型）
example_input = torch.randn(1, 10)
traced_model = torch.jit.trace(model, example_input)
traced_model.save('model_traced.pt')

# 加载并使用
loaded_model = torch.jit.load('model_scripted.pt')
with torch.no_grad():
    result = loaded_model(torch.randn(1, 10))
    print(f"输出形状: {result.shape}")

# setup.py 用于编译CUDA扩展
from setuptools import setup
from torch.utils.cpp_extension import BuildExtension, CUDAExtension

setup(
    name='custom_ops',
    ext_modules=[
        CUDAExtension('custom_ops', [
            'custom_ops.cpp',
            'custom_ops_cuda.cu',
        ])
    ],
    cmdclass={
        'build_ext': BuildExtension
    }
)

// custom_ops_cuda.cu - 自定义CUDA核函数示例
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>

__global__ void relu_kernel(float* input, float* output, int n) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < n) {
        output[idx] = input[idx] > 0 ? input[idx] : 0;
    }
}

torch::Tensor custom_relu_cuda(torch::Tensor input) {
    auto output = torch::zeros_like(input);
    int n = input.numel();
    
    const int threads = 256;
    const int blocks = (n + threads - 1) / threads;
    
    relu_kernel<<<blocks, threads>>>(
        input.data_ptr<float>(),
        output.data_ptr<float>(),
        n
    );
    
    return output;
}

import torch
import torchvision.models as models
import torchvision.transforms as transforms
from PIL import Image

# 加载预训练模型进行迁移学习
def create_transfer_model(num_classes, backbone='resnet50'):
    if backbone == 'resnet50':
        model = models.resnet50(weights=models.ResNet50_Weights.IMAGENET1K_V2)
        
        # 冻结骨干网络参数
        for param in model.parameters():
            param.requires_grad = False
        
        # 替换分类头
        in_features = model.fc.in_features
        model.fc = torch.nn.Sequential(
            torch.nn.Dropout(0.5),
            torch.nn.Linear(in_features, 256),
            torch.nn.ReLU(),
            torch.nn.Dropout(0.3),
            torch.nn.Linear(256, num_classes)
        )
    
    elif backbone == 'vit_b_16':
        model = models.vit_b_16(weights=models.ViT_B_16_Weights.IMAGENET1K_V1)
        in_features = model.heads.head.in_features
        model.heads.head = torch.nn.Linear(in_features, num_classes)
    
    return model

# 图像预处理
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]
    )
])

# 单张图片推理
def predict_image(model, image_path, class_names, device):
    model.eval()
    image = Image.open(image_path).convert('RGB')
    input_tensor = transform(image).unsqueeze(0).to(device)
    
    with torch.no_grad():
        outputs = model(input_tensor)
        probabilities = torch.softmax(outputs, dim=1)
        top5_prob, top5_idx = torch.topk(probabilities, 5)
    
    print("Top-5 预测结果:")
    for prob, idx in zip(top5_prob[0], top5_idx[0]):
        print(f"  {class_names[idx]}: {prob.item()*100:.2f}%")

import torch
import torch.nn as nn
from transformers import BertTokenizer, BertModel

class BertClassifier(nn.Module):
    """基于BERT的文本分类模型"""
    
    def __init__(self, num_classes, dropout=0.3):
        super(BertClassifier, self).__init__()
        self.bert = BertModel.from_pretrained('bert-base-chinese')
        self.dropout = nn.Dropout(dropout)
        self.classifier = nn.Linear(768, num_classes)
    
    def forward(self, input_ids, attention_mask, token_type_ids=None):
        outputs = self.bert(
            input_ids=input_ids,
            attention_mask=attention_mask,
            token_type_ids=token_type_ids
        )
        
        # 使用[CLS]标记的表示
        pooled_output = outputs.pooler_output
        pooled_output = self.dropout(pooled_output)
        logits = self.classifier(pooled_output)
        return logits


def prepare_text_data(texts, labels, tokenizer, max_length=128):
    """文本数据预处理"""
    encodings = tokenizer(
        texts,
        truncation=True,
        padding=True,
        max_length=max_length,
        return_tensors='pt'
    )
    return encodings, torch.tensor(labels)


# 使用示例
tokenizer = BertTokenizer.from_pretrained('bert-base-chinese')
model = BertClassifier(num_classes=2)  # 二分类（情感分析）

texts = ["这部电影真的非常精彩！", "服务态度很差，非常失望。"]
labels = [1, 0]  # 1=正面, 0=负面

encodings, label_tensor = prepare_text_data(texts, labels, tokenizer)

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = model.to(device)

with torch.no_grad():
    outputs = model(
        encodings['input_ids'].to(device),
        encodings['attention_mask'].to(device),
        encodings.get('token_type_ids', None)
    )
    predictions = torch.argmax(outputs, dim=1)
    print(f"预测结果: {predictions.cpu().numpy()}")

import torch
import torch.nn as nn
import torch.optim as optim
import numpy as np
from collections import deque
import random

class DQNNetwork(nn.Module):
    """深度Q网络"""
    
    def __init__(self, state_dim, action_dim, hidden_dim=256):
        super(DQNNetwork, self).__init__()
        self.network = nn.Sequential(
            nn.Linear(state_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, action_dim)
        )
    
    def forward(self, x):
        return self.network(x)


class ReplayBuffer:
    """经验回放缓冲区"""
    
    def __init__(self, capacity=10000):
        self.buffer = deque(maxlen=capacity)
    
    def push(self, state, action, reward, next_state, done):
        self.buffer.append((state, action, reward, next_state, done))
    
    def sample(self, batch_size):
        batch = random.sample(self.buffer, batch_size)
        states, actions, rewards, next_states, dones = zip(*batch)
        return (np.array(states), np.array(actions),
                np.array(rewards), np.array(next_states), np.array(dones))
    
    def __len__(self):
        return len(self.buffer)


class DQNAgent:
    """DQN智能体"""
    
    def __init__(self, state_dim, action_dim, lr=1e-3, gamma=0.99):
        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        self.action_dim = action_dim
        self.gamma = gamma
        
        # 在线网络和目标网络
        self.q_network = DQNNetwork(state_dim, action_dim).to(self.device)
        self.target_network = DQNNetwork(state_dim, action_dim).to(self.device)
        self.target_network.load_state_dict(self.q_network.state_dict())
        
        self.optimizer = optim.Adam(self.q_network.parameters(), lr=lr)
        self.replay_buffer = ReplayBuffer()
        self.epsilon = 1.0
        self.epsilon_decay = 0.995
        self.epsilon_min = 0.01
    
    def select_action(self, state):
        if random.random() < self.epsilon:
            return random.randint(0, self.action_dim - 1)
        
        with torch.no_grad():
            state_tensor = torch.FloatTensor(state).unsqueeze(0).to(self.device)
            q_values = self.q_network(state_tensor)
            return q_values.argmax().item()
    
    def update(self, batch_size=64):
        if len(self.replay_buffer) < batch_size:
            return
        
        states, actions, rewards, next_states, dones = \
            self.replay_buffer.sample(batch_size)
        
        states = torch.FloatTensor(states).to(self.device)
        actions = torch.LongTensor(actions).to(self.device)
        rewards = torch.FloatTensor(rewards).to(self.device)
        next_states = torch.FloatTensor(next_states).to(self.device)
        dones = torch.FloatTensor(dones).to(self.device)
        
        # 计算当前Q值
        current_q = self.q_network(states).gather(1, actions.unsqueeze(1))
        
        # 计算目标Q值
        with torch.no_grad():
            next_q = self.target_network(next_states).max(1)[0]
            target_q = rewards + self.gamma * next_q * (1 - dones)
        
        # 计算Huber损失
        loss = nn.SmoothL1Loss()(current_q.squeeze(), target_q)
        
        self.optimizer.zero_grad()
        loss.backward()
        self.optimizer.step()
        
        # 衰减探索率
        self.epsilon = max(self.epsilon_min, self.epsilon * self.epsilon_decay)
        
        return loss.item()

import pytorch_lightning as pl
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader, random_split
from torchvision import datasets, transforms

class LightningCNN(pl.LightningModule):
    """使用PyTorch Lightning构建的CNN分类器"""
    
    def __init__(self, num_classes=10, learning_rate=1e-3):
        super().__init__()
        self.save_hyperparameters()
        
        self.conv_layers = nn.Sequential(
            nn.Conv2d(1, 32, 3, padding=1),
            nn.BatchNorm2d(32),
            nn.ReLU(),
            nn.MaxPool2d(2),
            nn.Conv2d(32, 64, 3, padding=1),
            nn.BatchNorm2d(64),
            nn.ReLU(),
            nn.MaxPool2d(2)
        )
        self.fc_layers = nn.Sequential(
            nn.Flatten(),
            nn.Linear(64 * 7 * 7, 128),
            nn.ReLU(),
            nn.Dropout(0.5),
            nn.Linear(128, num_classes)
        )
    
    def forward(self, x):
        return self.fc_layers(self.conv_layers(x))
    
    def training_step(self, batch, batch_idx):
        x, y = batch
        logits = self(x)
        loss = F.cross_entropy(logits, y)
        acc = (logits.argmax(dim=1) == y).float().mean()
        
        self.log('train_loss', loss, prog_bar=True)
        self.log('train_acc', acc, prog_bar=True)
        return loss
    
    def validation_step(self, batch, batch_idx):
        x, y = batch
        logits = self(x)
        loss = F.cross_entropy(logits, y)
        acc = (logits.argmax(dim=1) == y).float().mean()
        
        self.log('val_loss', loss, prog_bar=True)
        self.log('val_acc', acc, prog_bar=True)
    
    def configure_optimizers(self):
        optimizer = torch.optim.AdamW(
            self.parameters(), 
            lr=self.hparams.learning_rate,
            weight_decay=0.01
        )
        scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
            optimizer, mode='min', patience=5
        )
        return {
            'optimizer': optimizer,
            'lr_scheduler': {'scheduler': scheduler, 'monitor': 'val_loss'}
        }


# 数据模块
class MNISTDataModule(pl.LightningDataModule):
    def __init__(self, data_dir='./data', batch_size=64):
        super().__init__()
        self.data_dir = data_dir
        self.batch_size = batch_size
        self.transform = transforms.Compose([
            transforms.ToTensor(),
            transforms.Normalize((0.1307,), (0.3081,))
        ])
    
    def setup(self, stage=None):
        mnist_full = datasets.MNIST(
            self.data_dir, train=True, transform=self.transform, download=True
        )
        self.mnist_train, self.mnist_val = random_split(mnist_full, [55000, 5000])
        self.mnist_test = datasets.MNIST(
            self.data_dir, train=False, transform=self.transform, download=True
        )
    
    def train_dataloader(self):
        return DataLoader(self.mnist_train, batch_size=self.batch_size, shuffle=True)
    
    def val_dataloader(self):
        return DataLoader(self.mnist_val, batch_size=self.batch_size)
    
    def test_dataloader(self):
        return DataLoader(self.mnist_test, batch_size=self.batch_size)


# 训练
model = LightningCNN()
datamodule = MNISTDataModule()

trainer = pl.Trainer(
    max_epochs=20,
    accelerator='auto',      # 自动选择CPU/GPU
    devices='auto',
    precision=16,            # 混合精度
    callbacks=[
        pl.callbacks.EarlyStopping(monitor='val_loss', patience=5),
        pl.callbacks.ModelCheckpoint(
            monitor='val_acc', 
            mode='max',
            save_top_k=3
        ),
        pl.callbacks.LearningRateMonitor()
    ],
    logger=pl.loggers.TensorBoardLogger('logs/', name='mnist_cnn')
)

trainer.fit(model, datamodule)
trainer.test(model, datamodule)

# 准备模型存档
# mar_file_name.mar 的创建命令：
# torch-model-archiver \
#   --model-name my_model \
#   --version 1.0 \
#   --model-file model.py \
#   --serialized-file model.pth \
#   --handler image_classifier \
#   --extra-files index_to_name.json

# 自定义处理器
from ts.torch_handler.base_handler import BaseHandler
import torch
import json

class CustomHandler(BaseHandler):
    def preprocess(self, data):
        """数据预处理"""
        inputs = []
        for row in data:
            input_data = row.get('data') or row.get('body')
            # 处理输入数据
            tensor = self._preprocess_input(input_data)
            inputs.append(tensor)
        return torch.stack(inputs)
    
    def inference(self, data):
        """模型推理"""
        with torch.no_grad():
            return self.model(data)
    
    def postprocess(self, data):
        """结果后处理"""
        results = []
        for output in data:
            probs = torch.softmax(output, dim=0)
            top5_prob, top5_idx = torch.topk(probs, 5)
            result = {
                str(idx.item()): prob.item() 
                for prob, idx in zip(top5_prob, top5_idx)
            }
            results.append(result)
        return results

import torch
import torch.nn as nn
from torch.profiler import profile, record_function, ProfilerActivity

model = nn.Sequential(
    nn.Linear(1000, 500),
    nn.ReLU(),
    nn.Linear(500, 100)
).cuda()

inputs = torch.randn(64, 1000).cuda()

# 使用PyTorch Profiler
with profile(
    activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
    record_shapes=True,
    profile_memory=True,
    with_stack=True
) as prof:
    with record_function("model_inference"):
        for _ in range(50):
            output = model(inputs)

# 打印性能报告
print(prof.key_averages().table(
    sort_by="cuda_time_total", 
    row_limit=10
))

# 导出Chrome Trace格式，可在chrome://tracing中可视化
prof.export_chrome_trace("trace.json")

# 导出到TensorBoard
prof.export_stacks("profiler_stacks.txt", "self_cuda_time_total")

import torch
import gc

def optimize_memory():
    """内存优化常用技巧"""
    
    # 1. 及时清理不需要的张量
    def clear_cache():
        gc.collect()
        torch.cuda.empty_cache()
    
    # 2. 使用梯度检查点（以计算换内存）
    from torch.utils.checkpoint import checkpoint
    
    class MemoryEfficientModel(nn.Module):
        def __init__(self):
            super().__init__()
            self.layers = nn.ModuleList([
                nn.Linear(512, 512) for _ in range(10)
            ])
        
        def forward(self, x):
            for layer in self.layers:
                # 使用checkpoint节省显存
                x = checkpoint(layer, x)
            return x
    
    # 3. 使用in-place操作
    x = torch.randn(1000, 1000)
    x.relu_()  # in-place版本，节省内存
    
    # 4. 监控GPU内存使用
    def print_gpu_memory():
        if torch.cuda.is_available():
            print(f"GPU内存已分配: {torch.cuda.memory_allocated()/1024**2:.1f} MB")
            print(f"GPU内存已缓存: {torch.cuda.memory_reserved()/1024**2:.1f} MB")
            print(f"GPU内存最大使用: {torch.cuda.max_memory_allocated()/1024**2:.1f} MB")
    
    return print_gpu_memory

# 解决方案：梯度裁剪 + 适当的初始化
model = MyModel()

# Xavier初始化
for m in model.modules():
    if isinstance(m, nn.Linear):
        nn.init.xavier_uniform_(m.weight)
        nn.init.constant_(m.bias, 0)
    elif isinstance(m, nn.Conv2d):
        nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')

# 训练时进行梯度裁剪
optimizer.zero_grad()
loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
optimizer.step()

# 监控梯度范数
for name, param in model.named_parameters():
    if param.grad is not None:
        grad_norm = param.grad.norm().item()
        if grad_norm > 10:
            print(f"警告：{name} 梯度范数过大: {grad_norm:.4f}")

# 多种正则化技术组合使用
class RegularizedModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.layers = nn.Sequential(
            nn.Linear(100, 256),
            nn.BatchNorm1d(256),   # 批归一化
            nn.ReLU(),
            nn.Dropout(0.4),       # Dropout
            nn.Linear(256, 128),
            nn.LayerNorm(128),     # 层归一化
            nn.ReLU(),
            nn.Dropout(0.3),
            nn.Linear(128, 10)
        )
    
    def forward(self, x):
        return self.layers(x)

# L2正则化通过优化器的weight_decay参数实现
optimizer = torch.optim.AdamW(
    model.parameters(),
    lr=1e-3,
    weight_decay=1e-4  # L2正则化系数
)

import torch
import torch.nn as nn

# torch.compile的高级用法
model = nn.Transformer(d_model=512, nhead=8)

# 自定义编译后端
@torch.compile(backend='inductor', dynamic=True)
def optimized_forward(model, src, tgt):
    return model(src, tgt)

# dynamic=True支持动态形状，适合序列长度可变的场景
src = torch.randn(10, 32, 512)
tgt = torch.randn(20, 32, 512)
output = optimized_forward(model, src, tgt)

import torch
import torch.nn.functional as F

# FlashAttention集成（需要安装flash-attn库）
# PyTorch 2.0+ 内置了scaled_dot_product_attention的优化实现
def efficient_attention(query, key, value, mask=None):
    """使用PyTorch内置的高效注意力机制"""
    with torch.backends.cuda.sdp_kernel(
        enable_flash=True,
        enable_math=False,
        enable_mem_efficient=True
    ):
        output = F.scaled_dot_product_attention(
            query, key, value,
            attn_mask=mask,
            dropout_p=0.1 if training else 0.0,
            is_causal=False
        )
    return output

import torch
import torch.quantization as quantization

# 动态量化（适用于LSTM、Linear层）
model = MyLSTMModel()
quantized_model = quantization.quantize_dynamic(
    model,
    {torch.nn.LSTM, torch.nn.Linear},
    dtype=torch.qint8
)

# 静态量化（需要校准数据）
def quantize_static(model, calibration_loader):
    model.eval()
    
    # 融合层（提高量化效率）
    model_fused = quantization.fuse_modules(
        model, 
        [['conv', 'bn', 'relu']]
    )
    
    # 配置量化方案
    model_fused.qconfig = quantization.get_default_qconfig('x86')
    
    # 插入量化/反量化节点
    model_prepared = quantization.prepare(model_fused)
    
    # 使用校准数据运行前向传播
    with torch.no_grad():
        for inputs, _ in calibration_loader:
            model_prepared(inputs)
    
    # 转换为量化模型
    quantized_model = quantization.convert(model_prepared)
    
    return quantized_model

import torch

# 为移动端优化模型
model.eval()

# 使用torch.jit.script生成移动端兼容的模型
scripted_model = torch.jit.script(model)

# 针对移动端优化
from torch.utils.mobile_optimizer import optimize_for_mobile
optimized_model = optimize_for_mobile(scripted_model)

# 保存为移动端格式
optimized_model._save_for_lite_interpreter("model_mobile.ptl")
print("移动端模型已保存")

# 模型大小对比
import os
original_size = os.path.getsize("model.pt") / 1024 / 1024
mobile_size = os.path.getsize("model_mobile.ptl") / 1024 / 1024
print(f"原始模型: {original_size:.2f} MB")
print(f"移动端模型: {mobile_size:.2f} MB")