[AI学习笔记]稀疏化训练：DeepSeek万亿参数管理秘籍

# 示例代码：动态掩码生成器
class DynamicMaskGenerator(nn.Module):
    def __init__(self, density=0.3):
        super().__init__()
        self.density = density

    def forward(self, weight):
        mask = torch.rand_like(weight) 
        threshold = torch.quantile(mask, self.density)
        return (mask > threshold).float()

pip install torch torchvision torchaudio

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

class SimpleCNN(nn.Module):
    def __init__(self):
        super(SimpleCNN, self).__init__()
        self.conv1 = nn.Conv2d(3, 32, 3, padding=1)
        self.conv2 = nn.Conv2d(32, 64, 3, padding=1)
        self.fc1 = nn.Linear(64 * 8 * 8, 512)
        self.fc2 = nn.Linear(512, 10)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        x = F.max_pool2d(x, 2)
        x = F.relu(self.conv2(x))
        x = F.max_pool2d(x, 2)
        x = x.view(x.size(0), -1)
        x = F.relu(self.fc1(x))
        x = self.fc2(x)
        return x

from torchvision import datasets, transforms
from torch.utils.data import DataLoader

# 定义数据预处理变换
transform = transforms.Compose([
    transforms.RandomHorizontalFlip(),
    transforms.RandomCrop(32, padding=4),
    transforms.ToTensor(),
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])

# 加载训练集和测试集
train_dataset = datasets.CIFAR10(root='./data', train=True, download=True, transform=transform)
test_dataset = datasets.CIFAR10(root='./data', train=False, download=True, transform=transform)

# 创建数据加载器
train_loader = DataLoader(train_dataset, batch_size=128, shuffle=True, num_workers=4)
test_loader = DataLoader(test_dataset, batch_size=128, shuffle=False, num_workers=4)

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = SimpleCNN().to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
criterion = nn.CrossEntropyLoss()

# 稀疏化训练参数
sparsity_level = 0.3  # 稀疏度水平，表示希望30%的参数变为零
prune_every = 10      # 每10个批次进行一次剪枝

def train_sparse(model, train_loader, criterion, optimizer, sparsity_level, prune_every):
    model.train()
    for batch_idx, (data, target) in enumerate(train_loader):
        data, target = data.to(device), target.to(device)
        optimizer.zero_grad()
        output = model(data)
        loss = criterion(output, target)
        loss.backward()
        optimizer.step()

        # 每prune_every个批次进行一次剪枝
        if (batch_idx + 1) % prune_every == 0:
            # 获取模型参数
            parameters = list(model.parameters())
            # 计算所有参数的绝对值
            abs_params = [param.abs() for param in parameters]
            # 将所有参数展平为一维
            all_abs_params = torch.cat([param.view(-1) for param in abs_params])
            # 计算需要保留的参数数量
            k = int((1 - sparsity_level) * all_abs_params.numel())
            # 获取第k小的值作为阈值
            threshold = torch.kthvalue(all_abs_params, k).values
            # 对每个参数应用剪枝
            for param in parameters:
                mask = param.abs() >= threshold
                param.data = param.data * mask.float()

        if (batch_idx + 1) % 100 == 0:
            print(f'Train Batch: {batch_idx + 1}/{len(train_loader)} Loss: {loss.item():.4f}')

# 开始稀疏化训练
for epoch in range(50):
    print(f'Epoch: {epoch + 1}')
    train_sparse(model, train_loader, criterion, optimizer, sparsity_level, prune_every)

def evaluate(model, test_loader):
    model.eval()
    correct = 0
    total = 0
    with torch.no_grad():
        for data, target in test_loader:
            data, target = data.to(device), target.to(device)
            output = model(data)
            _, predicted = torch.max(output.data, 1)
            total += target.size(0)
            correct += (predicted == target).sum().item()
    accuracy = 100 * correct / total
    print(f'Test Accuracy: {accuracy:.2f}%')
    return accuracy

# 评估模型
test_accuracy = evaluate(model, test_loader)

# 保存模型
torch.save(model.state_dict(), 'sparse_model.pth')
print('Model saved successfully!')