import torch
import torch.nn as nn
import random
-------------------------
超参数
-------------------------
N_max = 10 # 最大 n 值（训练时用）
input_size = 4 # 字符集大小：'0', '1',
hidden_size = 64
output_size = input_size
num_layers = 1
lr = 0.01
epochs = 2000
-------------------------
1. 生成训练数据
-------------------------
def generate_sequences(N_max):
seqs = []
for n in range(1, N_max + 1):
s = '0' * n + '1' * n
seqs.append(s)
return seqs
sequences = generate_sequences(N_max)
字符到索引的映射
char_to_idx = {'0': 0, '1': 1, '
idx_to_char = {v: k for k, v in char_to_idx.items()}
vocab_size = len(char_to_idx)
将字符串转换为索引序列（带和）
def string_to_tensor(s):
tensor = [char_to_idx['
return torch.tensor(tensor, dtype=torch.long)
所有序列转为张量
data = [string_to_tensor(s) for s in sequences]
-------------------------
2. 定义 RNN 模型
-------------------------
class SimpleRNN(nn.Module):
def init(self, input_size, hidden_size, output_size, num_layers=1):
super(SimpleRNN, self).init()
self.hidden_size = hidden_size
self.num_layers = num_layers
self.embedding = nn.Embedding(input_size, hidden_size)
self.rnn = nn.GRU(hidden_size, hidden_size, num_layers, batch_first=True)
self.fc = nn.Linear(hidden_size, output_size)
def forward(self, x, hidden=None):embedded = self.embedding(x)output, hidden = self.rnn(embedded, hidden)logits = self.fc(output)return logits, hidden
model = SimpleRNN(vocab_size, hidden_size, vocab_size, num_layers)
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=lr)
-------------------------
3. 训练模型
-------------------------
print("开始训练...")
for epoch in range(epochs):
total_loss = 0
for seq in data:
optimizer.zero_grad()
# 输入：去掉最后一个 token（
inputs = seq[:-1].unsqueeze(0) # (1, L)
# 目标：去掉第一个 token（
targets = seq[1:].unsqueeze(0) # (1, L)
logits, _ = model(inputs)# logits: (1, L, vocab_size) → reshapeloss = criterion(logits.view(-1, vocab_size), targets.view(-1))loss.backward()optimizer.step()total_loss += loss.item()
if (epoch + 1) % 500 == 0:print(f"Epoch {epoch+1}, Loss: {total_loss:.4f}")
print("训练完成！")
-------------------------
4. 生成指定 n 的序列
-------------------------
def generate_sequence(n):
model.eval()
with torch.no_grad():
input_token = torch.tensor([[char_to_idx['
hidden = None
generated = []
for _ in range(2 * n + 1): # 最多生成 2n+1 个 token（含
logits, hidden = model(input_token, hidden)
prob = torch.softmax(logits[0, -1], dim=0)
next_token = torch.multinomial(prob, 1).item()
if next_token == char_to_idx['
break
generated.append(idx_to_char[next_token])
input_token = torch.tensor([[next_token]], dtype=torch.long)
return ''.join(generated)
用户输入 n
try:
n = int(input("请输入 n："))
if n < 1:
raise ValueError
result = generate_sequence(n)
print(f"生成的序列 (n={n}): {result}")
# 验证
expected = '0' * n + '1' * n
print(f"期望序列: {expected}")
print("✅ 正确!" if result == expected else "❌ 错误!")
except Exception as e:
print("输入无效，请输入正整数。")
