1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
| import torch
import torch.nn as nn
import time
class SequenceDataset:
"""简单的序列分类数据集"""
def __init__(self, seq_len=50, vocab_size=100, num_classes=2, n_samples=1000):
self.seq_len = seq_len
self.vocab_size = vocab_size
self.num_classes = num_classes
self.n_samples = n_samples
self.X = torch.randint(0, vocab_size, (n_samples, seq_len))
self.y = (self.X[:, :10].sum() > self.X[:, -10:].sum()).long()
self.y = torch.where(self.y, torch.ones(n_samples), torch.zeros(n_samples)).long()
def get_dataloader(self, batch_size=32, shuffle=True):
from torch.utils.data import TensorDataset, DataLoader
dataset = TensorDataset(self.X.float(), self.y)
return DataLoader(dataset, batch_size=batch_size, shuffle=shuffle)
def benchmark_models():
"""对比三种模型的训练速度和性能"""
torch.manual_seed(42)
seq_len = 50
vocab_size = 100
hidden_size = 64
embedding_dim = 32
num_classes = 2
batch_size = 64
train_data = SequenceDataset(seq_len, vocab_size, num_classes, 2000)
test_data = SequenceDataset(seq_len, vocab_size, num_classes, 500)
train_loader = train_data.get_dataloader(batch_size, shuffle=True)
test_loader = test_data.get_dataloader(batch_size, shuffle=False)
models = {
'RNN': nn.RNN(embedding_dim, hidden_size, batch_first=True),
'LSTM': nn.LSTM(embedding_dim, hidden_size, batch_first=True),
'GRU': nn.GRU(embedding_dim, hidden_size, batch_first=True),
}
results = {}
for name, encoder in models.items():
print(f"\n{'='*50}")
print(f"训练 {name}")
print('='*50)
model = nn.Sequential(
nn.Embedding(vocab_size, embedding_dim),
encoder,
nn.Linear(hidden_size, num_classes)
)
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
criterion = nn.CrossEntropyLoss()
start = time.time()
model.train()
for epoch in range(10):
total_loss = 0
for xb, yb in train_loader:
optimizer.zero_grad()
emb = model[0](xb.long())
out, _ = model[1](emb)
logits = model[2](out[:, -1, :])
loss = criterion(logits, yb)
loss.backward()
optimizer.step()
total_loss += loss.item()
if (epoch + 1) % 5 == 0:
print(f" Epoch {epoch+1}/10, Loss: {total_loss/len(train_loader):.4f}")
train_time = time.time() - start
model.eval()
correct = 0
total = 0
with torch.no_grad():
for xb, yb in test_loader:
emb = model[0](xb.long())
out, _ = model[1](emb)
logits = model[2](out[:, -1, :])
preds = logits.argmax(dim=1)
correct += (preds == yb).sum().item()
total += len(yb)
accuracy = correct / total
results[name] = {'train_time': train_time, 'accuracy': accuracy}
print(f" 训练时间: {train_time:.2f}s")
print(f" 测试准确率: {accuracy:.2%}")
print(f"\n{'='*50}")
print("对比总结")
print('='*50)
for name, metrics in results.items():
print(f"{name:6s} | 训练时间: {metrics['train_time']:6.2f}s | 准确率: {metrics['accuracy']:.2%}")
return results
results = benchmark_models()
|
