搭建模型¶

class Perceptron(nn.Module):
    """
    Perceptron.

    Rosenblatt, F. (1958). The perceptron: a probabilistic model for information storage and organization in the brain. Psychological review, 65(6), 386.
    """

    def __init__(self, input_size, output_size=1, episilon=1e-6):
        super(Perceptron, self).__init__()
        self.linear = nn.Linear(input_size, output_size)
        self.episolon = episilon

    def forward(self, x):
        t = self.linear(x)
        return (t + self.episolon) / torch.abs(t + self.episolon)

class MLP(nn.Module):
    """
    Multilayer Perceptron.
    
    Rumelhart, D. E., Hinton, G. E., & Williams, R. J. (1986). Learning representations by back-propagating errors. nature, 323(6088), 533-536.
    """

    def __init__(self, input_size, hidden_size, output_size=1):
        super(MLP, self).__init__()
        self.hidden = nn.Linear(input_size, hidden_size)
        self.output = nn.Linear(hidden_size, output_size)

    def forward(self, x):
        x = torch.sigmoid(self.hidden(x))
        x = torch.sigmoid(self.output(x))
        return x

mlp = MLP(2, 5, 1)
loss_func = nn.BCELoss()
optimizer = torch.optim.Adam(mlp.parameters(), lr=0.1)
log = []
for epoch in range(200):
    total_loss = 0
    pred = mlp(X)
    loss = loss_func(pred, Y)
    mlp.zero_grad()
    loss.backward()
    optimizer.step()
    total_loss += loss.item()
    log.append(total_loss / 4)

class ElmanRNN(nn.Module):
    """
    Elman RNN.

    Elman, J. L. (1990). Finding structure in time. Cognitive science, 14(2), 179-211.
    """

    def __init__(self, input_size, hidden_size, output_size=1):
        super(ElmanRNN, self).__init__()
        self.linear_x = nn.Linear(input_size, hidden_size)
        self.linear_h = nn.Linear(hidden_size, hidden_size)
        self.output = nn.Linear(hidden_size, output_size)

    def forward(self, x, hidden_state=None):
        if hidden_state is None:
            # 初始化隐藏状态为零
            hidden_state = torch.zeros(x.size(0), self.linear_h.out_features)
        hidden_state = torch.tanh(self.linear_x(x) + self.linear_h(hidden_state))
        return torch.sigmoid(self.output(hidden_state)), hidden_state

rnn_cell = nn.RNNCell(3, 2)

input_size, hidden_size, num_layers = 3, 5, 2
rnn = nn.RNN(input_size, hidden_size, num_layers)
params = dict(rnn.named_parameters())
x = torch.randn(2, 4, 3)

def forward(x, hx=None, batch_first=False):
    if batch_first:
        x = x.transpose(0, 1)
    seq_len, batch_size, _ = x.size()
    if hx is None:
        hx = torch.zeros(rnn.num_layers, batch_size, rnn.hidden_size)
    h_t_minus_1 = hx.clone()
    h_t = hx.clone()
    output = []
    for t in range(seq_len):
        for layer in range(rnn.num_layers):
            input_t = x[t] if layer == 0 else h_t[layer - 1]
            h_t[layer] = torch.tanh(
                input_t @ params[f"weight_ih_l{layer}"].T
                + h_t_minus_1[layer] @ params[f"weight_hh_l{layer}"].T
                + params[f"bias_hh_l{layer}"]
                + params[f"bias_ih_l{layer}"]
            )
        output.append(h_t[-1].clone())
        h_t_minus_1 = h_t.clone()
    output = torch.stack(output)
    if batch_first:
        output = output.transpose(0, 1)
    return output, h_t

official_imp = rnn(x)
my_imp = forward(x)

# 由于浮点精度误差的存在，我们这里不使用torch.equal来判断两个tensor是否相等
assert torch.allclose(official_imp[0], my_imp[0])
assert torch.allclose(official_imp[1], my_imp[1])

class LeNet(nn.Module):
    """
    LeNet.

    LeCun, Y., Bottou, L., Bengio, Y., & Haffner, P. (1998). Gradient-based learning applied to document recognition. Proceedings of the IEEE, 86(11), 2278-2324.
    """

    def __init__(self, num_classes=10):
        super(LeNet, self).__init__()
        self.conv1 = nn.Conv2d(1, 6, kernel_size=5, stride=1, padding=2)
        self.conv2 = nn.Conv2d(6, 16, kernel_size=5, stride=1)
        self.fc1 = nn.Linear(16 * 5 * 5, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, num_classes)

    def forward(self, x):
        x = torch.tanh(self.conv1(x))
        x = nn.functional.avg_pool2d(x, kernel_size=2, stride=2)
        x = torch.tanh(self.conv2(x))
        x = nn.functional.avg_pool2d(x, kernel_size=2, stride=2)
        x = x.view(x.size(0), -1)  # Flatten the tensor
        x = torch.tanh(self.fc1(x))
        x = torch.tanh(self.fc2(x))
        x = self.fc3(x)
        return x

lenet = LeNet()

class AlexNet(nn.Module):
    """
    AlexNet.

    Krizhevsky, A., Sutskever, I., & Hinton, G. E. (2012). Imagenet classification with deep convolutional neural networks. Advances in neural information processing systems, 25.
    """
    def __init__(self, num_classes: int = 1000, dropout: float = 0.5) -> None:
        super().__init__()
        self.features = nn.Sequential(
            nn.Conv2d(3, 64, kernel_size=11, stride=4, padding=2),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=3, stride=2),
            nn.Conv2d(64, 192, kernel_size=5, padding=2),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=3, stride=2),
            nn.Conv2d(192, 384, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(384, 256, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(256, 256, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=3, stride=2),
        )
        self.avgpool = nn.AdaptiveAvgPool2d((6, 6))
        self.classifier = nn.Sequential(
            nn.Dropout(p=dropout),
            nn.Linear(256 * 6 * 6, 4096),
            nn.ReLU(inplace=True),
            nn.Dropout(p=dropout),
            nn.Linear(4096, 4096),
            nn.ReLU(inplace=True),
            nn.Linear(4096, num_classes),
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.features(x)
        x = self.avgpool(x)
        x = torch.flatten(x, 1)
        x = self.classifier(x)
        return x

deconv = nn.ConvTranspose2d(
    in_channels=1,        # 输入通道数
    out_channels=1,       # 输出通道数
    kernel_size=3,        # 卷积核大小
    stride=1,             # 步长
    padding=0,            # 填充
    output_padding=0,     # 输出填充
    dilation=1,           # 扩张（空洞卷积）
    groups=1,             # 分组卷积
    bias=True,            # 是否使用偏置
    padding_mode="zeros", # 填充模式
    device=None,          # 设备
    dtype=None,           # 数据类型
)

class DoubleConv(nn.Module):
    """(convolution => [BN] => ReLU) * 2"""

    def __init__(self, in_channels, out_channels, mid_channels=None):
        super().__init__()
        if not mid_channels:
            mid_channels = out_channels
        self.double_conv = nn.Sequential(
            nn.Conv2d(in_channels, mid_channels, kernel_size=3, padding=1, bias=False),
            nn.BatchNorm2d(mid_channels),
            nn.ReLU(inplace=True),
            nn.Conv2d(mid_channels, out_channels, kernel_size=3, padding=1, bias=False),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(inplace=True),
        )

    def forward(self, x):
        return self.double_conv(x)


class Down(nn.Module):
    """Downscaling with maxpool then double conv"""

    def __init__(self, in_channels, out_channels):
        super().__init__()
        self.maxpool_conv = nn.Sequential(
            nn.MaxPool2d(2), DoubleConv(in_channels, out_channels)
        )

    def forward(self, x):
        return self.maxpool_conv(x)


class Up(nn.Module):
    """Upscaling then double conv"""

    # 默认情况下是先做（双线性）上采样，然后再做卷积
    # 也可以是先做转置卷积，然后再做卷积

    def __init__(self, in_channels, out_channels, bilinear=True):
        super().__init__()

        # if bilinear, use the normal convolutions to reduce the number of channels
        if bilinear:
            self.up = nn.Upsample(scale_factor=2, mode="bilinear", align_corners=True)
            self.conv = DoubleConv(in_channels, out_channels, in_channels // 2)
        else:
            # 转置卷积！
            self.up = nn.ConvTranspose2d(
                in_channels, in_channels // 2, kernel_size=2, stride=2
            )
            self.conv = DoubleConv(in_channels, out_channels)

    def forward(self, x1, x2):
        x1 = self.up(x1)
        diffY = x2.size()[2] - x1.size()[2]
        diffX = x2.size()[3] - x1.size()[3]
        x1 = nn.functional.pad(
            x1, [diffX // 2, diffX - diffX // 2, diffY // 2, diffY - diffY // 2]
        )
        x = torch.cat([x2, x1], dim=1)
        return self.conv(x)


class UNet(nn.Module):
    """
    U-Net.

    Ronneberger, O., Fischer, P., & Brox, T. (2015). U-net: Convolutional networks for biomedical image segmentation. In Medical image computing and computer-assisted intervention–MICCAI 2015: 18th international conference, Munich, Germany, October 5-9, 2015, proceedings, part III 18 (pp. 234-241). Springer international publishing.
    """

    def __init__(self, n_channels, n_classes, bilinear=False):
        super(UNet, self).__init__()
        self.n_channels = n_channels
        self.n_classes = n_classes
        self.bilinear = bilinear

        self.inc = DoubleConv(n_channels, 64)
        self.down1 = Down(64, 128)
        self.down2 = Down(128, 256)
        self.down3 = Down(256, 512)
        factor = 2 if bilinear else 1
        self.down4 = Down(512, 1024 // factor)
        self.up1 = Up(1024, 512 // factor, bilinear)
        self.up2 = Up(512, 256 // factor, bilinear)
        self.up3 = Up(256, 128 // factor, bilinear)
        self.up4 = Up(128, 64, bilinear)
        self.outc = nn.Conv2d(64, n_classes, kernel_size=1)

    def forward(self, x):
        x1 = self.inc(x)
        x2 = self.down1(x1)
        x3 = self.down2(x2)
        x4 = self.down3(x3)
        x5 = self.down4(x4)
        x = self.up1(x5, x4)
        x = self.up2(x, x3)
        x = self.up3(x, x2)
        x = self.up4(x, x1)
        logits = self.outc(x)
        return logits

# 为了方便复现，我们设置drop概率为0
drop = nn.Dropout(0)

# 编码层自注意力
## QVK一起计算
encoder_projection = (
    src @ params["encoder.layers.0.self_attn.in_proj_weight"].T
    + params["encoder.layers.0.self_attn.in_proj_bias"]
)
## 分离计算QKV
q, k, v = encoder_projection.split(26, dim=-1)
## 需要转置为 (1,2,10,13) 也就是 (batch, head, seq_len, d_model/num_heads)
## 从而在后续的矩阵乘法，把13消掉
q = q.view(10, 1, 2, 13).permute(1, 2, 0, 3)
k = k.view(10, 1, 2, 13).permute(1, 2, 0, 3)
v = v.view(10, 1, 2, 13).permute(1, 2, 0, 3)
## 计算注意力分数，这是一个(seq_len, seq_len)的矩阵
## 因为我们这里计算的是 ** 一个句子内部，不同词之间的相关性 **
score = torch.softmax(q @ k.transpose(-1, -2) / (13**0.5), dim=-1)
## 计算注意力值
attention = score @ v
attention_out = (
    attention.transpose(1, 2).contiguous().view(10, 1, 26)
    @ params["encoder.layers.0.self_attn.out_proj.weight"].T
    + params["encoder.layers.0.self_attn.out_proj.bias"]
)
attention_out = drop(attention_out)
# add & norm
self_attention = torch.nn.functional.layer_norm(
    src + attention_out,
    [26],
    weight=params["encoder.layers.0.norm1.weight"],
    bias=params["encoder.layers.0.norm1.bias"],
)
# FFN
ffn_out = torch.relu(
    self_attention @ params["encoder.layers.0.linear1.weight"].T
    + params["encoder.layers.0.linear1.bias"]
)
ffn_out = drop(ffn_out)
ffn_out = (
    ffn_out @ params["encoder.layers.0.linear2.weight"].T
    + params["encoder.layers.0.linear2.bias"]
)
ffn_out = drop(ffn_out)
# add & norm
encoder_out = torch.nn.functional.layer_norm(
    self_attention + ffn_out,
    [26],
    weight=params["encoder.layers.0.norm2.weight"],
    bias=params["encoder.layers.0.norm2.bias"],
)
encoder_out = torch.nn.functional.layer_norm(
    encoder_out,
    [26],
    weight=params["encoder.norm.weight"],
    bias=params["encoder.norm.bias"],
)

# 解码层自注意力
## 和编码层自注意力完全一致，把输入替换为tgt即可
decoder_projection = (
    tgt @ params["decoder.layers.0.self_attn.in_proj_weight"].T
    + params["decoder.layers.0.self_attn.in_proj_bias"]
)
q, k, v = decoder_projection.split(26, dim=-1)
q = q.view(10, 1, 2, 13).permute(1, 2, 0, 3)
k = k.view(10, 1, 2, 13).permute(1, 2, 0, 3)
v = v.view(10, 1, 2, 13).permute(1, 2, 0, 3)
score = q @ k.transpose(-1, -2) / (13**0.5)
mask = torch.tril(torch.ones(1, 2, 10, 10))
score = score.masked_fill(mask == 0, -1e9)
score = torch.softmax(score, dim=-1)
attention = score @ v
out = (
    attention.transpose(1, 2).contiguous().view(10, 1, 26)
    @ params["decoder.layers.0.self_attn.out_proj.weight"].T
    + params["decoder.layers.0.self_attn.out_proj.bias"]
)
out = drop(out)
tgt_attention = torch.nn.functional.layer_norm(
    tgt + out,
    [26],
    weight=params["decoder.layers.0.norm1.weight"],
    bias=params["decoder.layers.0.norm1.bias"],
)

# 交叉注意力（cross attention）
wq, wk, wv = params["decoder.layers.0.multihead_attn.in_proj_weight"].split(26, dim=0)
bq, bk, bv = params["decoder.layers.0.multihead_attn.in_proj_bias"].split(26, dim=0)
proj_Q = tgt_attention @ wq.T + bq
proj_K = encoder_out @ wk.T + bk
proj_V = encoder_out @ wv.T + bv
q = proj_Q.view(10, 1, 2, 13).permute(1, 2, 0, 3)
k = proj_K.view(10, 1, 2, 13).permute(1, 2, 0, 3)
v = proj_V.view(10, 1, 2, 13).permute(1, 2, 0, 3)
score = torch.softmax(q @ k.transpose(-1, -2) / (13**0.5), dim=-1)
attention = score @ v
out = (
    attention.transpose(1, 2).contiguous().view(10, 1, 26)
    @ params["decoder.layers.0.multihead_attn.out_proj.weight"].T
    + params["decoder.layers.0.multihead_attn.out_proj.bias"]
)
out = drop(out)
# add & norm
cross_attention = torch.nn.functional.layer_norm(
    tgt_attention + out,
    [26],
    weight=params["decoder.layers.0.norm2.weight"],
    bias=params["decoder.layers.0.norm2.bias"],
)
# FFN
ffn_out = torch.relu(
    cross_attention @ params["decoder.layers.0.linear1.weight"].T
    + params["decoder.layers.0.linear1.bias"]
)
ffn_out = drop(ffn_out)
ffn_out = (
    ffn_out @ params["decoder.layers.0.linear2.weight"].T
    + params["decoder.layers.0.linear2.bias"]
)
ffn_out = drop(ffn_out)
# add & norm
decoder_out = torch.nn.functional.layer_norm(
    cross_attention + ffn_out,
    [26],
    weight=params["decoder.layers.0.norm3.weight"],
    bias=params["decoder.layers.0.norm3.bias"],
)
decoder_out = torch.nn.functional.layer_norm(
    decoder_out,
    [26],
    weight=params["decoder.norm.weight"],
    bias=params["decoder.norm.bias"],
)