LSTM 单元前向传播（AI033）- 算法机考高频题型 · 循环网络 / LSTM题解 | AlgoMooc

AI033. LSTM 单元前向传播

困难

通过率 29% · 提交 7 · 通过 2

循环网络LSTM

这道 OJ 的配套动画

LSTM 前向传播 · 对应本题的思路、过程和代码

本题改编自算法机考高频题型 · 考点：循环网络 / LSTM · 难度：困难

实现单层长短期记忆网络（LSTM）的前向传播（输入维 = 隐藏维 = d）：输入门 i=σ(...)、遗忘门 f=σ(...)、候选值 g=tanh(...)、输出门 o=σ(...)；细胞状态 c_t = f⊙c_{t−1} + i⊙g；隐藏态 h_t = o⊙tanh(c_t)。σ 为 sigmoid 函数，σ(z) = 1/(1+e^{−z})。h_0、c_0 由输入给出，输出最终的 h_T 与 c_T。

这类题属于算法机考高频题型中「循环网络 / LSTM」方向的高频题型，通常考察对「循环网络 / LSTM」的建模能力与边界条件处理。掌握本题的解题思路后，可举一反三应对同类真题方向，稳步提升机考通过率。

输入描述

第一行两个整数 T d。随后按 i、f、g、o 的固定顺序给出四个门的参数，每个门依次为 d×d 的 W_x、d×d 的 W_h、1 行 d 个数的 b；随后 1 行 h_0、1 行 c_0；随后 T 行、每行 d 个实数为 x_1…x_T。

输出描述

第一行 d 个数为 h_T，第二行 d 个数为 c_T（四舍五入保留两位小数，单个空格分隔）；-0.00 输出 0.00。舍入规则：用语言的标准两位小数格式化（f"{x:.2f}"、printf("%.2f")、String.format("%.2f")）即可；测试数据保证每个待输出值的精确值到 0.005 舍入边界的距离不小于 1e-8，三种语言结果一致，无需自行实现舍入。

数据范围

1 <= T <= 30；1 <= d <= 8；|权重、h_0、c_0、x| <= 4；|偏置| <= 10（用于构造门激活值接近 1 或接近 0 的情形，样例中的 10.00 即此用途）。

示例

示例 1

输入示例

输出示例

0.96
1.96

说明：

d=1，各门激活值接近 1

时间限制 4000 ms · 内存限制 256 MB

本平台为独立第三方培训机构，与华为技术有限公司无任何关联；课程的服务内容与权益以购买协议为准，学习效果因个人情况而异。「华为 OD」「华为可信」等仅为对岗位与考试方向的客观描述，相关商标归各自权利人所有。

问小欧 · 图解算法年卡¥99 开通 →

题库上一题下一题

未登录也能直接「运行」调试;「提交」判题需微信登录,你写的代码会自动保留、登录后回到本题继续。

字号13Ctrl+/ 注释 · Ctrl+↵ 提交

1
2

# 在此编写你的代码

写完代码点「提交」，将对全部测试用例判题。

解题思路

每个时间步用旧隐藏态分别计算 i、f、g、o 四个门，再先更新 c_t=f⊙c_{t-1}+i⊙g，最后计算 h_t=o⊙tanh(c_t)。

正确性说明

对时间步归纳：四个门均由同一个正确的 h_{t-1} 与当前输入得到；细胞状态严格组合旧状态和本步候选，隐藏态再读取本步新状态。该顺序与 LSTM 定义一致，所以最终 c_T、h_T 正确。

复杂度分析

四个门各做两次矩阵向量乘，时间 O(Td²)，滚动状态空间 O(d)。

边界与易错点

参数必须按 i、f、g、o 读取；遗忘门接近 1 或 0 时分别检查持续累积与清空旧状态，避免把新 h 提前用于同一步其他门。

参考代码

import math
import sys


def fmt(v: float) -> str:
    s = f"{v:.2f}"
    return "0.00" if s == "-0.00" else s


def sigmoid(z: float) -> float:
    if z >= 0:
        return 1.0 / (1.0 + math.exp(-z))
    e = math.exp(z)
    return e / (1.0 + e)


def main() -> None:
    data = sys.stdin.read().split()
    it = iter(data)
    T = int(next(it)); d = int(next(it))
    read_mat = lambda: [[float(next(it)) for _ in range(d)] for _ in range(d)]
    read_vec = lambda: [float(next(it)) for _ in range(d)]
    gates = []  # 顺序 i f g o：每门 (Wx, Wh, b)
    for _ in range(4):
        wx = read_mat(); wh = read_mat(); b = read_vec()
        gates.append((wx, wh, b))
    h = read_vec()
    c = read_vec()

    def affine(gate, x, hprev):
        wx, wh, b = gate
        out = []
        for r in range(d):
            z = b[r]
            z += sum(wx[r][t] * x[t] for t in range(d))
            z += sum(wh[r][t] * hprev[t] for t in range(d))
            out.append(z)
        return out

    for _ in range(T):
        x = read_vec()
        zi = affine(gates[0], x, h)
        zf = affine(gates[1], x, h)
        zg = affine(gates[2], x, h)
        zo = affine(gates[3], x, h)
        i_g = [sigmoid(z) for z in zi]
        f_g = [sigmoid(z) for z in zf]
        g_g = [math.tanh(z) for z in zg]
        o_g = [sigmoid(z) for z in zo]
        c = [f_g[r] * c[r] + i_g[r] * g_g[r] for r in range(d)]  # 用上一步 c
        h = [o_g[r] * math.tanh(c[r]) for r in range(d)]         # 用本步 c
    sys.stdout.write(" ".join(fmt(v) for v in h) + "\n")
    sys.stdout.write(" ".join(fmt(v) for v in c) + "\n")


if __name__ == "__main__":
    main()