Activations

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

# 1) ReLU — the workhorse
relu = nn.ReLU()
relu(torch.tensor([-2.0, 0.0, 3.0]))      # tensor([0., 0., 3.])

# Functional form (no parameters, no module)
F.relu(x)

# 2) Leaky ReLU + Parametric ReLU — fixes 'dying ReLU'
lrelu = nn.LeakyReLU(negative_slope=0.01) # leaks 1 percent of negative
prelu = nn.PReLU(num_parameters=1)        # learns the slope

# 3) GELU — used in Transformers (BERT, GPT)
gelu = nn.GELU()
F.gelu(x)
# Smoother than ReLU; better empirically for attention models.

# 4) SiLU / Swish — used in EfficientNet, Llama, modern CNNs
silu = nn.SiLU()                          # x * sigmoid(x)
F.silu(x)

# 5) Tanh — bounded (-1, 1)
F.tanh(x)
# Good for RNNs, regression outputs in a fixed range.
# Saturates → vanishing gradients in deep nets.

# 6) Sigmoid — bounded (0, 1)
F.sigmoid(x)
# Use only at output for binary classification or gating; don't stack in hidden layers.

# 7) Softmax — turns logits into probabilities
logits = torch.tensor([[2.0, 1.0, 0.1]])
probs  = F.softmax(logits, dim=-1)        # [0.659, 0.242, 0.099]
logp   = F.log_softmax(logits, dim=-1)    # numerically stabler

# Common bug: applying softmax then CrossEntropyLoss (which expects logits)
# → double softmax. Pass raw logits to nn.CrossEntropyLoss directly.

# 8) Where to place activations
class MLP(nn.Module):
    def __init__(self, d_in=784, d_hidden=256, d_out=10):
        super().__init__()
        self.fc1 = nn.Linear(d_in, d_hidden)
        self.fc2 = nn.Linear(d_hidden, d_hidden)
        self.fc3 = nn.Linear(d_hidden, d_out)

    def forward(self, x):
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        return self.fc3(x)                 # raw logits — no softmax

# 9) Custom activation
class Mish(nn.Module):
    def forward(self, x):
        return x * torch.tanh(F.softplus(x))

mish = Mish()

# 10) Gating — modern Transformer FFN (GLU variants)
class SwiGLU(nn.Module):
    def __init__(self, d_model, d_ff):
        super().__init__()
        self.w1 = nn.Linear(d_model, d_ff, bias=False)
        self.w2 = nn.Linear(d_model, d_ff, bias=False)
        self.w3 = nn.Linear(d_ff, d_model, bias=False)

    def forward(self, x):
        return self.w3(F.silu(self.w1(x)) * self.w2(x))

# 11) Initialization matters
linear = nn.Linear(256, 256)
nn.init.kaiming_normal_(linear.weight, nonlinearity='relu')   # He init for ReLU
nn.init.xavier_uniform_(linear.weight)                          # Xavier for tanh/sigmoid

# 12) Diagnostics — watch dead neurons
def relu_alive_fraction(model, x):
    activations = []
    def hook(m, i, o): activations.append((o > 0).float().mean().item())
    handles = [m.register_forward_hook(hook) for m in model.modules() if isinstance(m, nn.ReLU)]
    model(x); [h.remove() for h in handles]
    return activations
# < 0.1 alive fraction → many dead neurons → use LeakyReLU or lower LR

# 13) Quick reference — what to use
#   Hidden layers (general):  ReLU
#   Transformers FFN:          GELU or SwiGLU
#   CNNs (modern):             SiLU/Swish
#   RNN/LSTM gates:            sigmoid + tanh
#   Output: classification:    softmax (or none + CrossEntropyLoss)
#   Output: regression:         no activation (linear)
#   Output: bounded regression: tanh or sigmoid

# 14) Common bugs
#   • Softmax then CrossEntropyLoss → use raw logits
#   • ReLU on output of a regression model → can't predict negatives
#   • Sigmoid everywhere → vanishing gradients, slow training
#   • Forgetting dim= in softmax → softmaxes over the wrong axis
#   • Putting activation before, not after, batchnorm → known suboptimal

import torch.nn.functional as F
F.relu(x)
F.gelu(x)
F.silu(x)
F.softmax(x, dim=-1)