import torch
import torch.nn as nn
import torch.optim as optim
import numpy as np
====================== 1. 配置参数（初学者友好版） ======================
设备配置：优先使用GPU，没有则用CPU
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
序列参数：n的取值范围（生成0n1n，n从1到5）
n_list = [1, 2, 3, 4, 5]
字符到数字的映射（编码）
char2idx = {'0': 0, '1': 1}
idx2char = {0: '0', 1: '1'}
vocab_size = len(char2idx) # 词汇表大小（0和1）
RNN参数
input_size = vocab_size # 输入维度（one-hot编码）
hidden_size = 16 # 隐藏层维度（初学者可先固定）
num_layers = 1 # RNN层数（单层更简单）
learning_rate = 0.01 # 学习率
num_epochs = 1000 # 训练轮数
print_every = 200 # 每200轮打印一次训练结果
====================== 2. 生成训练数据 ======================
def generate_data(n_list):
"""生成0n1n形式的序列数据"""
data = []
for n in n_list:
# 生成序列：n个0 + n个1
seq = '0' * n + '1' * n
data.append(seq)
return data
生成训练数据
train_data = generate_data(n_list)
print("训练序列样本：", train_data)
====================== 3. 数据预处理（编码+转换为张量） ======================
def encode_sequence(seq):
"""将字符序列转换为数字编码的张量（one-hot）"""
# 第一步：字符转数字索引
idx_seq = [char2idx[char] for char in seq]
# 第二步：转换为one-hot编码（模型输入需要）
one_hot_seq = torch.zeros(len(seq), vocab_size).to(device)
for i, idx in enumerate(idx_seq):
one_hot_seq[i, idx] = 1.0
return one_hot_seq, torch.tensor(idx_seq).to(device)
====================== 4. 构建简单的RNN模型 ======================
class SimpleRNN(nn.Module):
def init(self, input_size, hidden_size, vocab_size, num_layers):
super(SimpleRNN, self).init()
self.hidden_size = hidden_size # 隐藏层维度
self.num_layers = num_layers # RNN层数
# 定义RNN层：输入维度=词汇表大小，隐藏层维度=hidden_size，层数=num_layers
self.rnn = nn.RNN(input_size, hidden_size, num_layers, batch_first=True)
# 全连接层：将隐藏层输出映射到词汇表大小（预测0或1）
self.fc = nn.Linear(hidden_size, vocab_size)
def forward(self, x, hidden):"""前向传播x: 输入张量 (batch_size, seq_len, input_size)hidden: 隐藏状态张量 (num_layers, batch_size, hidden_size)"""# RNN前向传播out, hidden = self.rnn(x, hidden)# 调整形状：(batch_size*seq_len, hidden_size)，方便全连接层处理out = out.reshape(out.size(0)*out.size(1), out.size(2))# 全连接层预测out = self.fc(out)return out, hiddendef init_hidden(self, batch_size):"""初始化隐藏状态（全0）"""return torch.zeros(self.num_layers, batch_size, self.hidden_size).to(device)
====================== 5. 初始化模型、损失函数、优化器 ======================
model = SimpleRNN(input_size, hidden_size, vocab_size, num_layers).to(device)
criterion = nn.CrossEntropyLoss() # 交叉熵损失（分类任务）
optimizer = optim.Adam(model.parameters(), lr=learning_rate) # Adam优化器
====================== 6. 训练模型 ======================
print("\n开始训练模型...")
for epoch in range(num_epochs):
total_loss = 0
# 遍历所有训练序列
for seq in train_data:
# 编码序列：获取one-hot输入和数字标签
input_one_hot, target_idx = encode_sequence(seq)
# 调整输入形状：(batch_size=1, seq_len, input_size)
input_one_hot = input_one_hot.unsqueeze(0)
# 初始化隐藏状态hidden = model.init_hidden(batch_size=1)# 梯度清零（避免累积）optimizer.zero_grad()# 前向传播：逐个字符预测（除了最后一个字符）loss = 0for i in range(len(seq)-1):# 取第i个字符作为输入，预测第i+1个字符input_char = input_one_hot[:, i:i+1, :]output, hidden = model(input_char, hidden)# 计算损失：预测值 vs 第i+1个字符的标签loss += criterion(output, target_idx[i+1:i+2])# 反向传播+参数更新loss.backward()optimizer.step()total_loss += loss.item()# 打印训练进度
if (epoch+1) % print_every == 0:avg_loss = total_loss / len(train_data)print(f'Epoch [{epoch+1}/{num_epochs}], 平均损失: {avg_loss:.4f}')
====================== 7. 测试模型：生成0^n 1^n序列 ======================
print("\n测试模型生成序列...")
def generate_sequence(n):
"""
生成0^n 1^n序列
n: 0和1的个数
"""
# 初始输入：第一个字符是0（one-hot编码）
input_char = torch.tensor([[char2idx['0']]]).to(device)
input_one_hot = torch.zeros(1, 1, vocab_size).to(device)
input_one_hot[0, 0, input_char] = 1.0
# 初始化隐藏状态
hidden = model.init_hidden(batch_size=1)
# 存储生成的序列
generated_seq = ['0']# 第一步：生成前n-1个0
for _ in range(n-1):output, hidden = model(input_one_hot, hidden)# 预测下一个字符（取概率最大的）_, predicted_idx = torch.max(output, 1)generated_seq.append(idx2char[predicted_idx.item()])# 更新输入为当前预测的字符input_one_hot = torch.zeros(1, 1, vocab_size).to(device)input_one_hot[0, 0, predicted_idx] = 1.0# 第二步：生成n个1
for _ in range(n):output, hidden = model(input_one_hot, hidden)_, predicted_idx = torch.max(output, 1)generated_seq.append(idx2char[predicted_idx.item()])return ''.join(generated_seq)
测试生成不同n的序列
for test_n in [1, 2, 3, 4, 5]:
generated_seq = generate_sequence(test_n)
expected_seq = '0'test_n + '1'test_n
print(f'n={test_n}: 预期序列={expected_seq}, 生成序列={generated_seq}, 匹配: {generated_seq == expected_seq}')