P128

import torch
import torch.nn as nn
import torch.optim as optim
import random

1. 生成训练数据

def generate_data(max_n=5, num_samples=1000):
data = []
for _ in range(num_samples):
n = random.randint(1, max_n)
seq = [0]n + [1]n
data.append(seq)
return data

2. 定义LSTM模型

class LSTMGenerator(nn.Module):
def init(self, input_size=2, hidden_size=32, output_size=2):
super().init()
self.lstm = nn.LSTM(input_size, hidden_size, batch_first=True)
self.fc = nn.Linear(hidden_size, output_size)

def forward(self, x, hidden=None):
    out, hidden = self.lstm(x, hidden)
    out = self.fc(out)
    return out, hidden

3. 数据编码（0→[1,0]，1→[0,1]）

def encode_seq(seq):
return torch.tensor([[1,0] if c == 0 else [0,1] for c in seq], dtype=torch.float32).unsqueeze(0)

4. 训练模型

data = generate_data()
model = LSTMGenerator()
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

for epoch in range(1000):
total_loss = 0
for seq in data:
x = encode_seq(seq[:-1]) # 输入前n-1个字符
y = torch.tensor(seq[1:], dtype=torch.long) # 预测后n-1个字符
optimizer.zero_grad()
out, _ = model(x)
loss = criterion(out.squeeze(0), y)
loss.backward()
optimizer.step()
total_loss += loss.item()
if (epoch+1) % 100 == 0:
print(f"Epoch {epoch+1}, Loss: {total_loss/len(data):.4f}")

5. 生成序列（比如生成n=3的000111）

def generate_seq(model, n):
model.eval()
seq = [0]*n
with torch.no_grad():
# 先输入n个0
x = encode_seq(seq)
out, hidden = model(x)
# 生成n个1
for _ in range(n):
next_char = torch.argmax(out[:, -1, :]).item()
seq.append(next_char)
# 用最后一个输出作为下一个输入
next_input = encode_seq([next_char]).squeeze(0)
out, hidden = model(next_input.unsqueeze(0), hidden)
return seq

测试生成

generated = generate_seq(model, 3)
print("生成的序列:", generated) # 应该输出[0,0,0,1,1,1]