brain_ai_train.py

"""
BNA Training on MNIST - 可运行的类脑AI训练
运行: python brain_ai_train.py
"""

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.utils.data as data
import torchvision
import numpy as np
import math
import time

# ================== 核心模块 ==================

class SpikingLIF(nn.Module):
    """脉冲LIF神经元"""
    def __init__(self, dim: int, tau: float = 2.0):
        super().__init__()
        self.fc = nn.Linear(dim, dim)
        self.tau = tau
        self.v_mem = None

    def forward(self, x):
        if self.v_mem is None or self.v_mem.shape != x.shape:
            self.v_mem = torch.zeros_like(x)
        h = self.fc(x)
        self.v_mem = self.v_mem * math.exp(-1.0 / self.tau) + h
        spikes = (self.v_mem > 1.0).float()
        self.v_mem = self.v_mem * (1 - spikes)
        return spikes + (h - h.detach())


class EventAttention(nn.Module):
    """稀疏事件注意力"""
    def __init__(self, dim: int, heads: int = 8, k_ratio: float = 0.15):
        super().__init__()
        self.heads = heads
        self.k_ratio = k_ratio
        self.head_dim = dim // heads
        self.qkv = nn.Linear(dim, dim * 3)
        self.proj = nn.Linear(dim, dim)

    def forward(self, x):
        B, N, C = x.shape
        qkv = self.qkv(x).reshape(B, N, 3, self.heads, self.head_dim).permute(2, 0, 3, 1, 4)
        q, k, v = qkv[0], qkv[1], qkv[2]
        attn = (q @ k.transpose(-2, -1)) / math.sqrt(self.head_dim)
        k = max(1, int(N * self.k_ratio))
        top_val, top_idx = torch.topk(attn, k=k, dim=-1)
        mask = torch.full_like(attn, float('-inf')).scatter_(-1, top_idx, top_val)
        attn = F.softmax(mask, dim=-1)
        out = (attn @ v).transpose(1, 2).reshape(B, N, C)
        return self.proj(out)


class BNA(nn.Module):
    """类脑神经网络"""
    def __init__(self):
        super().__init__()
        self.input_proj = nn.Linear(784, 256)

        self.spike1 = SpikingLIF(256, tau=2.0)
        self.spike2 = SpikingLIF(256, tau=2.0)

        self.predictive1 = nn.Sequential(
            nn.Linear(256, 512), nn.GELU(), nn.Linear(512, 256)
        )
        self.predictive2 = nn.Sequential(
            nn.Linear(256, 512), nn.GELU(), nn.Linear(512, 256)
        )

        self.attn = EventAttention(256, heads=8)

        self.output = nn.Linear(256, 10)

    def forward(self, x):
        x = x.view(x.size(0), -1)
        x = self.input_proj(x)

        x = self.spike1(x)
        err1 = x - self.predictive1(x)
        x = x + err1 * 0.1

        x = self.spike2(x)
        err2 = x - self.predictive2(x)
        x = x + err2 * 0.1

        x = x.unsqueeze(1)
        x = self.attn(x)
        x = x.squeeze(1)

        return self.output(x)


class Transformer(nn.Module):
    """标准Transformer对比"""
    def __init__(self):
        super().__init__()
        self.input_proj = nn.Linear(784, 256)
        self.cls_token = nn.Parameter(torch.zeros(1, 1, 256))
        self.encoder = nn.TransformerEncoder(
            nn.TransformerEncoderLayer(d_model=256, nhead=8, dim_feedforward=1024, batch_first=True),
            num_layers=3
        )
        self.output = nn.Linear(256, 10)

    def forward(self, x):
        x = x.view(x.size(0), -1)
        x = self.input_proj(x).unsqueeze(1)
        cls = self.cls_token.expand(x.size(0), -1, -1)
        x = torch.cat([cls, x], dim=1)
        x = self.encoder(x)
        return self.output(x[:, 0])


def train_model(model, train_loader, test_loader, epochs=5, lr=1e-3, name="Model"):
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    model = model.to(device)
    opt = torch.optim.AdamW(model.parameters(), lr=lr)
    criterion = nn.CrossEntropyLoss()

    for epoch in range(epochs):
        model.train()
        total_loss = 0
        correct = 0
        total = 0
        start = time.time()

        for i, (img, label) in enumerate(train_loader):
            img, label = img.to(device), label.to(device)
            opt.zero_grad()
            out = model(img)
            loss = criterion(out, label)
            loss.backward()
            opt.step()

            total_loss += loss.item()
            correct += (out.argmax(1) == label).sum().item()
            total += label.size(0)

            if i % 100 == 0:
                print(f"  [{i}/{len(train_loader)}] Loss: {loss.item():.4f}")

        train_acc = correct / total
        train_time = time.time() - start

        # 测试
        model.eval()
        correct = 0
        total = 0
        with torch.no_grad():
            for img, label in test_loader:
                img, label = img.to(device), label.to(device)
                out = model(img)
                correct += (out.argmax(1) == label).sum().item()
                total += label.size(0)
        test_acc = correct / total

        print(f"{name} Epoch {epoch+1}/{epochs}: Train={train_acc:.4f}, Test={test_acc:.4f}, Time={train_time:.1f}s")


def main():
    print("="*60)
    print("BNA vs Transformer MNIST Training")
    print("="*60)

    # 数据
    transform = torchvision.transforms.Compose([
        torchvision.transforms.ToTensor(),
        lambda x: x.view(-1)
    ])
    train_data = torchvision.datasets.MNIST('./data', train=True, download=True, transform=transform)
    test_data = torchvision.datasets.MNIST('./data', train=False, download=True, transform=transform)
    train_loader = data.DataLoader(train_data, batch_size=128, shuffle=True)
    test_loader = data.DataLoader(test_data, batch_size=256)

    print(f"Train: {len(train_data)}, Test: {len(test_data)}")

    # 训练BNA
    print("\n--- Training BNA ---")
    bna = BNA()
    print(f"BNA Parameters: {sum(p.numel() for p in bna.parameters()):,}")
    train_model(bna, train_loader, test_loader, epochs=5, lr=1e-3, name="BNA")

    # 训练Transformer对比
    print("\n--- Training Transformer (Comparison) ---")
    transformer = Transformer()
    print(f"Transformer Parameters: {sum(p.numel() for p in transformer.parameters()):,}")
    train_model(transformer, train_loader, test_loader, epochs=5, lr=1e-3, name="Transformer")

    print("\nDone!")


if __name__ == "__main__":
    main()