長短期記憶網路(Long Short-Term Memory，LSTM)是一種專為處理序列資料所設計的遞迴神經網路(Recurrent Neural Network，RNN)，並被廣泛應用於時間序列預測、語音辨識、自然語言處理等需要捕捉時序依賴的任務。標準 RNN 在理論上可以記住任意長度的歷史資訊，但在實際訓練時，梯度會隨著時間步的增加而指數性地縮小或放大，導致難以學習長距離的依賴關係，此現象稱為梯度消失(vanishing gradient)或梯度爆炸(exploding gradient)問題。LSTM 透過引入細胞狀態(cell state)與三個閘門(gate)的設計，讓網路可以選擇性地記憶或遺忘資訊，從而有效緩解上述問題。

在每個時間步 t，LSTM 接收當前輸入 xt 與上一步的狀態 ht−1、ct−1，並計算出 ht 以及 ct，即：

(ht, ct) = LSTM(xt, ht-1, ct-1)

而其內部之詳細運算步驟，可分解如下：

ft = σ(Wf[ht-1, xt] + bf)

it = σ(Wi[ht-1, xt] + bi)

ot = σ(Wo[ht-1, xt] + bo)

c̃t = tanh(Wc[ht-1, xt] + bc)

ct = ft⊙ct-1 + it⊙c̃t

ht = ot⊙tanh(ct)

其中 W* 與 b* 為對應的權重與 bias，方括號表示矩陣並排，σ 為 sigmoid function，⊙ 為 element-wise product；ct-1 代表上一步的記憶，乘上 ft 代表上一步的記憶要保留多少；c̃t 代表這一步新產生的候選記憶，乘上 it 代表它要貢獻多少到這一步的新記憶當中。h0 與 c0 實務上通常使用 zero vector。

上述的詳細運算步驟，若要用 PyTorch class 的方式表示，則寫法如下：

class LSTMCell(nn.Module):
	def __init__(self, input_size, hidden_size):
		super().__init__()
		self.W_f = nn.Linear(input_size + hidden_size, hidden_size)
		self.W_i = nn.Linear(input_size + hidden_size, hidden_size)
		self.W_o = nn.Linear(input_size + hidden_size, hidden_size)
		self.W_c = nn.Linear(input_size + hidden_size, hidden_size)

	def forward(self, x, h_prev, c_prev):
		xh = torch.cat([h_prev, x], dim=-1)
		f = torch.sigmoid(self.W_f(xh))
		i = torch.sigmoid(self.W_i(xh))
		o = torch.sigmoid(self.W_o(xh))
		c_tilde = torch.tanh(self.W_c(xh))
		c = f * c_prev + i * c_tilde
		h = o * torch.tanh(c)
		return h, c

當然，實務上我們不必自己用自己撰寫的 LSTM cell 來搭建網路，而是使用 PyTorch 的 nn.LSTM。一個以 sine wave 做示範的簡單範例如下：

import matplotlib.pyplot as plt
import numpy as np
import torch
import torch.nn as nn
from torch.utils.data import DataLoader, TensorDataset

device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')

t_train = np.linspace(0, 100, 10000)
wave_train = np.sin(t_train).astype(np.float32)

seq_len = 50
X_train = np.array([wave_train[i:i + seq_len] for i in range(len(wave_train) - seq_len)])
y_train = wave_train[seq_len:]

X_train = torch.tensor(X_train).unsqueeze(-1) # (N, seq_len, 1)
y_train = torch.tensor(y_train).unsqueeze(-1) # (N, 1)

print('Data shapes:', X_train.shape, y_train.shape)

train_set = TensorDataset(X_train, y_train)
train_loader = DataLoader(train_set, batch_size=128, shuffle=True)


class LSTMModel(nn.Module):
	def __init__(self):
		super().__init__()
		self.lstm = nn.LSTM(input_size=1, hidden_size=64, batch_first=True)
		self.linear = nn.Linear(64, 1)

	def forward(self, x):
		out, _ = self.lstm(x)
		return self.linear(out[:, -1, :])


model = LSTMModel().to(device)
optim = torch.optim.Adam(model.parameters())
criterion = nn.MSELoss()

model.train()
for i in range(20):
	print(f'Epoch {i+1}/20')
	for X_batch, y_batch in train_loader:
		loss = criterion(model(X_batch.to(device)), y_batch.to(device))
		optim.zero_grad()
		loss.backward()
		optim.step()

t_end = 100 + 2 * np.pi
t_step = t_train[1] - t_train[0]
t_test = np.arange(100 + t_step, t_end + t_step, t_step, dtype=np.float32)

window = wave_train[-seq_len:].tolist()
preds = []

model.eval()
with torch.no_grad():
	for t in t_test:
		x_in = torch.tensor(window[-seq_len:], dtype=torch.float32).unsqueeze(0).unsqueeze(-1).to(device)
		pred = model(x_in).item()
		preds.append(pred)
		window.append(pred)

y_true = np.sin(t_test)

plt.plot(t_test, y_true, label='Ground Truth')
plt.plot(t_test, preds, label='Prediction', linestyle='--')
plt.xlabel('t')
plt.ylabel('sin(t)')
plt.title('LSTM Rolling Prediction')
plt.legend()
plt.show()

在上述範例中：

雖然 LSTM 實際上的應用，主要是跟時間序列有關的問題，但我們也可以將影像的其中一個軸視為時間軸，來套用 LSTM。以下是用 LSTM 來進行 MNIST 數字分類的範例：

import torch
import torch.nn as nn
from torch.utils.data import DataLoader
from torchvision import datasets, transforms

device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')

train_set = datasets.MNIST(root='./data', train=True,  download=True, transform=transforms.ToTensor())
test_set = datasets.MNIST(root='./data', train=False, download=True, transform=transforms.ToTensor())
print('Data shapes:', train_set.data.shape, test_set.data.shape)

train_loader = DataLoader(train_set, batch_size=128, shuffle=True)
test_loader = DataLoader(test_set, batch_size=128)


class LSTMModel(nn.Module):
	def __init__(self):
		super().__init__()
		self.lstm = nn.LSTM(input_size=28, hidden_size=64, batch_first=True)
		self.linear = nn.Linear(64, 10)

	def forward(self, x):
		x = x.squeeze(1)  # (N, 1, 28, 28) -> (N, 28, 28)
		out, _ = self.lstm(x)
		return self.linear(out[:, -1, :])


model = LSTMModel().to(device)
optim = torch.optim.Adam(model.parameters())
criterion = nn.CrossEntropyLoss()

model.train()
for i in range(10):
	print(f'Epoch {i+1}/10')
	for X_batch, y_batch in train_loader:
		loss = criterion(model(X_batch.to(device)), y_batch.to(device))
		optim.zero_grad()
		loss.backward()
		optim.step()

model.eval()
correct = sum(
	(model(X.to(device)).argmax(dim=1) == y.to(device)).sum().item()
	for X, y in test_loader
)
print('Accuracy: {:.2f}%'.format(100 * correct / len(test_set)))

在上述範例中，原始影像的 shape 是 (N, C=1, H=28, W=28)，我們用 squeeze(1) 來拿掉 channel 軸，並將 height 軸視為時間軸，餵給 LSTM。

既然 LSTM 的後面可以疊其他東西，那前面也可以。一個常見的作法是，把圖片等資料先用 CNN 處理過後，再餵給 LSTM 處理。若以基於上一個範例做修改來說，你只要把 class 替換為以下的版本即可：

class CNNLSTMModel(nn.Module):
	def __init__(self):
		super().__init__()
		self.cnn = nn.Sequential(
			nn.Conv2d(1, 16, kernel_size=3, padding=1),
			nn.ReLU(),
			nn.MaxPool2d(2),
			nn.Conv2d(16, 32, kernel_size=3, padding=1),
			nn.ReLU(),
			nn.MaxPool2d(2)
		)
		self.lstm = nn.LSTM(input_size=32*7, hidden_size=64, batch_first=True)
		self.linear = nn.Linear(64, 10)

	def forward(self, x):
		x = self.cnn(x) # (N, ch=1, h=28, w=28) to (N, ch=32, h=7, w=7)
		x = x.permute(0, 3, 1, 2) # (N, ch=32, h=7, w=7) to (N, w=7, ch=32, h=7)
		x = x.flatten(2) # (N, w=7, ch=32, h=7) to (N, 7, 32*7)
		out, _ = self.lstm(x)
		return self.linear(out[:, -1, :])

在上述的 class 中，你需要把 CNN 處理完的結果中，適合當時間軸的維度移到前面來，才能餵給 LSTM。此外，LSTM 的輸入大小必須事先計算，因此你需要相當確定原始的輸入大小。

關於前述內容中尚未示範的變體，包含但不限於以下：