AI模型框架实战指南：从CNN到Transformer的演进

class TextCNN(nn.Module):
    def __init__(self, vocab_size, embedding_dim, num_classes):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, embedding_dim)
        self.conv = nn.Conv1d(embedding_dim, 100, kernel_size=3)
        self.fc = nn.Linear(100, num_classes)

    def forward(self, x):
        x = self.embedding(x)
        x = x.transpose(1, 2)
        x = torch.relu(self.conv(x))
        x = torch.max_pool1d(x, x.size(2))
        x = x.squeeze(2)
        return self.fc(x)

from torchvision.models import resnet18

network = resnet18(weights=None)

# 1. 单通道输入（灰度图）
network.conv1 = nn.Conv2d(1, 64, kernel_size=5, stride=2, padding=2, bias=False)

# 2. 去掉 maxpool（小图保留更多分辨率）
network.maxpool = nn.Identity()

# 3. 去掉分类头，输出 embedding
features = network.fc.in_features
network.fc = nn.Identity()

self.backbone = network
self.embedding = nn.Sequential(
    nn.Dropout(0.15),
    nn.Linear(features, 256, bias=False),
    nn.BatchNorm1d(256),
)

# 4. 余弦分类头（比 softmax 更适合跨字体泛化）
self.classifier = CosineClassifier(256, num_classes, scale=30.0)

class SimpleRNN(nn.Module):
    def __init__(self, vocab_size, embedding_dim, hidden_dim):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, embedding_dim)
        self.rnn = nn.RNN(embedding_dim, hidden_dim, batch_first=True)

class MultiHeadAttention(nn.Module):
    def __init__(self, d_model, num_heads):
        super().__init__()
        assert d_model % num_heads == 0

        self.d_model = d_model
        self.num_heads = num_heads
        self.d_k = d_model // num_heads

        self.W_q = nn.Linear(d_model, d_model)
        self.W_k = nn.Linear(d_model, d_model)
        self.W_v = nn.Linear(d_model, d_model)
        self.W_o = nn.Linear(d_model, d_model)

    def scaled_dot_product_attention(self, Q, K, V, mask=None):
        attn_scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.d_k)

        if mask is not None:
            attn_scores = attn_scores.masked_fill(mask == 0, -1e9)

        attn_probs = torch.softmax(attn_scores, dim=-1)
        output = torch.matmul(attn_probs, V)
        return output, attn_probs

    def forward(self, x, mask=None):
        batch_size = x.size(0)

        Q = self.W_q(x).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2)
        K = self.W_k(x).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2)
        V = self.W_v(x).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2)

        attn_output, attn_probs = self.scaled_dot_product_attention(Q, K, V, mask)
        attn_output = attn_output.transpose(1, 2).contiguous().view(batch_size, -1, self.d_model)
        return self.W_o(attn_output)

class TransformerBlock(nn.Module):
    def __init__(self, embed_size, heads, dropout, forward_expansion):
        super().__init__()
        self.attention = MultiHeadAttention(embed_size, heads)
        self.norm1 = nn.LayerNorm(embed_size)
        self.norm2 = nn.LayerNorm(embed_size)

        self.feed_forward = nn.Sequential(
            nn.Linear(embed_size, forward_expansion * embed_size),
            nn.ReLU(),
            nn.Linear(forward_expansion * embed_size, embed_size)
        )
        self.dropout = nn.Dropout(dropout)

    def forward(self, value, key, query, mask):
        attention = self.attention(value, key, query, mask)
        x = self.dropout(self.norm1(attention + query))  # Add & Norm
        forward = self.feed_forward(x)
        out = self.dropout(self.norm2(forward + x))      # Add & Norm
        return out

class PositionalEncoding(nn.Module):
    def __init__(self, d_model, max_len=5000):
        super().__init__()
        position = torch.arange(0, max_len).unsqueeze(1).float()
        div_term = torch.exp(
            torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model)
        )

        pe = torch.zeros(max_len, d_model)
        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        pe = pe.unsqueeze(0)
        self.register_buffer('pe', pe)

    def forward(self, x):
        return x + self.pe[:, :x.size(1)]

def make_src_mask(src, pad_idx):
    # src: [B, L]，只要 != pad_idx 的位置
    return (src != pad_idx).unsqueeze(1).unsqueeze(2)  # [B, 1, 1, L]

def create_trg_mask(trg):
    N, trg_len = trg.shape
    trg_mask = torch.tril(torch.ones((trg_len, trg_len))).expand(
        N, 1, trg_len, trg_len
    )
    return trg_mask

def relu(x):
    return max(0, x)

import torch.nn.functional as F

# 方法 1：Swish 在 PyTorch 里叫 SiLU（beta=1 时等价）
output = F.silu(x)

# 方法 2：手动实现
def swish(x, beta=1.0):
    return x * torch.sigmoid(beta * x)

def safe_swish(x, beta=1.0):
    x = torch.clamp(x, -50.0, 50.0)
    return x * torch.sigmoid(beta * x)

# 训练循环中加梯度裁剪
loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
optimizer.step()

def get_linear_schedule_with_warmup(optimizer, num_warmup_steps, num_training_steps):
    def lr_lambda(current_step):
        if current_step < num_warmup_steps:
            return float(current_step) / float(max(1, num_warmup_steps))
        return max(0.0, float(num_training_steps - current_step) / 
                   float(max(1, num_training_steps - num_warmup_steps)))
    return LambdaLR(optimizer, lr_lambda)

from torch.nn.utils.rnn import pad_sequence

def collate_fn(batch):
    src_batch = [item['src'] for item in batch]
    trg_batch = [item['trg'] for item in batch]
    src_batch = pad_sequence(src_batch, batch_first=True, padding_idx=PAD_IDX)
    trg_batch = pad_sequence(trg_batch, batch_first=True, padding_idx=PAD_IDX)
    return {'src': src_batch, 'trg': trg_batch}

from torch.cuda.amp import autocast, GradScaler

scaler = GradScaler()
for batch in train_loader:
    optimizer.zero_grad()
    with autocast():
        output = model(src, trg[:-1, :])
        loss = criterion(output.reshape(-1, output.shape[2]), trg[1:, :].reshape(-1))
    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()

accumulation_steps = 4
for i, batch in enumerate(train_loader):
    with autocast():
        output = model(src)
        loss = criterion(output, target) / accumulation_steps
    scaler.scale(loss).backward()

    if (i + 1) % accumulation_steps == 0:
        scaler.step(optimizer)
        scaler.update()
        optimizer.zero_grad()

class ViTPatchEmbedding(nn.Module):
    def __init__(self, img_size=224, patch_size=16, in_channels=3, embed_dim=768):
        super().__init__()
        self.n_patches = (img_size // patch_size) ** 2
        # 本质上是一个 kernel=patch_size, stride=patch_size 的卷积
        self.proj = nn.Conv2d(in_channels, embed_dim,
                              kernel_size=patch_size, stride=patch_size)

    def forward(self, x):
        x = self.proj(x)         # [B, embed_dim, H/patch, W/patch]
        x = x.flatten(2)         # [B, embed_dim, n_patches]
        x = x.transpose(1, 2)    # [B, n_patches, embed_dim]
        return x

class VisionTransformer(nn.Module):
    def __init__(self, img_size=224, patch_size=16, in_channels=3, num_classes=1000,
                 embed_dim=768, depth=12, num_heads=12, mlp_ratio=4.0, dropout=0.1):
        super().__init__()
        self.patch_embed = ViTPatchEmbedding(img_size, patch_size, in_channels, embed_dim)
        self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
        self.pos_embed = nn.Parameter(torch.zeros(1, 1 + self.patch_embed.n_patches, embed_dim))
        self.pos_drop = nn.Dropout(dropout)

        self.blocks = nn.ModuleList([
            TransformerBlock(embed_dim, num_heads, dropout, mlp_ratio)
            for _ in range(depth)
        ])

        self.norm = nn.LayerNorm(embed_dim)
        self.head = nn.Linear(embed_dim, num_classes)

    def forward(self, x):
        x = self.patch_embed(x)
        cls_token = self.cls_token.expand(x.shape[0], -1, -1)
        x = torch.cat((cls_token, x), dim=1)
        x = x + self.pos_embed
        x = self.pos_drop(x)

        for block in self.blocks:
            x = block(x, x, x, None)

        x = self.norm(x)
        return self.head(x[:, 0])  # 用 CLS token 分类

def resize_pos_embed(posemb, new_grid_size):
    cls_token_emb = posemb[:, :1]
    grid_emb = posemb[:, 1:].reshape(1, int(math.sqrt(posemb.shape[1] - 1)),
                                      int(math.sqrt(posemb.shape[1] - 1)), -1)
    grid_emb = grid_emb.permute(0, 3, 1, 2)
    grid_emb = torch.nn.functional.interpolate(
        grid_emb, size=new_grid_size, mode='bicubic', align_corners=False)
    grid_emb = grid_emb.permute(0, 2, 3, 1).reshape(1, -1, grid_emb.shape[1])
    return torch.cat([cls_token_emb, grid_emb], dim=1)

# 推荐：AdamW + cosine schedule
optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4, weight_decay=0.05)
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=300)

from transformers import BertTokenizer, BertForSequenceClassification

tokenizer = BertTokenizer.from_pretrained('bert-base-chinese')
model = BertForSequenceClassification.from_pretrained('bert-base-chinese', num_labels=3)

text = "这家餐厅的菜品很不错，但服务有点慢"
inputs = tokenizer(text, return_tensors='pt', truncation=True, max_length=128)

with torch.no_grad():
    outputs = model(**inputs)
    predictions = torch.nn.functional.softmax(outputs.logits, dim=-1)

# 冻结前面层，只训练最后几层
for name, param in model.named_parameters():
    if "encoder.layer.0" in name or "encoder.layer.1" in name:
        param.requires_grad = False

# 或者用 LoRA / Adapter
optimizer = AdamW(filter(lambda p: p.requires_grad, model.parameters()), lr=2e-5)

def bert_masking(tokens, mask_prob=0.15):
    num_masks = int(len(tokens) * mask_prob)
    mask_indices = random.sample(range(len(tokens)), num_masks)

    masked_tokens = tokens.copy()
    for i in mask_indices:
        prob = random.random()
        if prob < 0.8:
            masked_tokens[i] = "[MASK]"                    # 80% 替换
        elif prob < 0.9:
            masked_tokens[i] = random.choice(all_tokens)   # 10% 随机词
        # 10% 保持原样
    return masked_tokens, mask_indices

def dynamic_masking(tokens, pos_tags, mask_prob=0.15):
    pos_weights = {'NOUN': 1.5, 'VERB': 1.5, 'ADJ': 1.2,
                   'ADV': 1.2, 'PRON': 0.8, 'PART': 0.5, 'default': 1.0}
    mask_indices = [i for i, tag in enumerate(pos_tags)
                    if random.random() < mask_prob * pos_weights.get(tag, 1.0)]

    masked_tokens = tokens.copy()
    for i in mask_indices:
        masked_tokens[i] = "[MASK]"
    return masked_tokens, mask_indices

# 错误：DataLoader 里动态掩码，多进程下模式不一致
dataloader = DataLoader(dataset, batch_size=32, num_workers=4,
                       collate_fn=lambda batch: dynamic_masking(batch))

# 正确：预处理阶段就掩码好
preprocessed_data = [dynamic_masking(t) for t in dataset]
dataloader = DataLoader(preprocessed_data, batch_size=32, num_workers=4)