Residual Connections

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

class BasicResidualBlock(nn.Module):
    """
    ResNet basic block: Conv → BN → ReLU → Conv → BN → (+input) → ReLU

    Nếu số kênh đầu vào và ra khác nhau hoặc stride > 1,
    dùng projection shortcut (1x1 conv) để khớp dimension.
    """

    expansion = 1

    def __init__(self, in_ch: int, out_ch: int, stride: int = 1):
        super().__init__()
        self.conv1 = nn.Conv2d(in_ch, out_ch, kernel_size=3,
                               stride=stride, padding=1, bias=False)
        self.bn1 = nn.BatchNorm2d(out_ch)
        self.conv2 = nn.Conv2d(out_ch, out_ch, kernel_size=3,
                               stride=1, padding=1, bias=False)
        self.bn2 = nn.BatchNorm2d(out_ch)

        # Skip connection: identity nếu dimension khớp,
        # còn không thì dùng 1x1 conv (projection shortcut).
        if stride != 1 or in_ch != out_ch * self.expansion:
            self.shortcut = nn.Sequential(
                nn.Conv2d(in_ch, out_ch * self.expansion,
                          kernel_size=1, stride=stride, bias=False),
                nn.BatchNorm2d(out_ch * self.expansion),
            )
        else:
            self.shortcut = nn.Identity()

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        identity = self.shortcut(x)

        # Nhánh chính F(x): Conv → BN → ReLU → Conv → BN
        out = self.conv1(x)
        out = self.bn1(out)
        out = F.relu(out, inplace=True)
        out = self.conv2(out)
        out = self.bn2(out)

        # Cộng với skip TRƯỚC khi qua ReLU cuối
        out = out + identity
        out = F.relu(out, inplace=True)
        return out


# Ví dụ: xếp chồng 3 block tạo thành một stage của ResNet
layers = nn.Sequential(
    BasicResidualBlock(in_ch=64, out_ch=64),
    BasicResidualBlock(in_ch=64, out_ch=64),
    BasicResidualBlock(in_ch=64, out_ch=128, stride=2),  # downsample
)
x = torch.randn(8, 64, 56, 56)
y = layers(x)
print(y.shape)  # torch.Size([8, 128, 28, 28])

import torch
import torch.nn as nn

class PreNormTransformerBlock(nn.Module):
    """
    Một block Transformer kiểu Pre-Norm (GPT-2, LLaMA, Mistral).

    2 residual connections mỗi block:
      1) x = x + MultiHeadAttention(LayerNorm(x))
      2) x = x + FFN(LayerNorm(x))

    Gradient chảy thẳng qua 2 đường skip không bị LayerNorm chặn.
    """

    def __init__(self, d_model: int, n_heads: int, d_ff: int,
                 dropout: float = 0.1):
        super().__init__()
        self.norm1 = nn.LayerNorm(d_model)
        self.attn = nn.MultiheadAttention(
            d_model, n_heads, dropout=dropout, batch_first=True
        )
        self.norm2 = nn.LayerNorm(d_model)
        self.ffn = nn.Sequential(
            nn.Linear(d_model, d_ff),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(d_ff, d_model),
        )
        self.drop = nn.Dropout(dropout)

    def forward(self, x: torch.Tensor,
                attn_mask: torch.Tensor | None = None) -> torch.Tensor:
        # Residual 1: quanh self-attention
        a = self.norm1(x)
        attn_out, _ = self.attn(a, a, a, attn_mask=attn_mask,
                                need_weights=False)
        x = x + self.drop(attn_out)

        # Residual 2: quanh FFN
        x = x + self.drop(self.ffn(self.norm2(x)))
        return x


# Xếp chồng 48 block như GPT-2 Medium
blocks = nn.Sequential(*[
    PreNormTransformerBlock(d_model=1024, n_heads=16, d_ff=4096)
    for _ in range(48)
])
x = torch.randn(2, 256, 1024)
y = blocks(x)
print(y.shape)  # torch.Size([2, 256, 1024])