129 lines
3.6 KiB
Python
129 lines
3.6 KiB
Python
#!/usr/bin/env python3
|
||
"""
|
||
快速TPU测试脚本 - 验证简单模型是否可以在TPU上运行
|
||
"""
|
||
|
||
import os
|
||
import time
|
||
import torch
|
||
import torch.nn as nn
|
||
|
||
# 设置环境变量
|
||
os.environ['XLA_FLAGS'] = '--xla_cpu_multi_thread_eigen=true --xla_cpu_enable_fast_math=true'
|
||
os.environ['XLA_USE_BF16'] = '1'
|
||
|
||
import torch_xla.core.xla_model as xm
|
||
|
||
def quick_test():
|
||
"""快速测试TPU是否工作正常"""
|
||
print("🚀 开始快速TPU测试...")
|
||
|
||
try:
|
||
# 获取TPU设备
|
||
device = xm.xla_device()
|
||
print(f"📱 TPU设备: {device}")
|
||
|
||
# 创建简单模型
|
||
model = nn.Sequential(
|
||
nn.Linear(512, 256),
|
||
nn.ReLU(),
|
||
nn.GRU(256, 128, batch_first=True),
|
||
nn.Linear(128, 41)
|
||
).to(device)
|
||
|
||
print(f"📊 模型参数: {sum(p.numel() for p in model.parameters()):,}")
|
||
|
||
# 创建测试数据
|
||
x = torch.randn(8, 50, 512, device=device)
|
||
print(f"📥 输入形状: {x.shape}")
|
||
|
||
# 测试前向传播
|
||
print("🔄 测试前向传播...")
|
||
start_time = time.time()
|
||
|
||
with torch.no_grad():
|
||
if hasattr(model, '__getitem__'):
|
||
# 对于Sequential模型，手动处理GRU层
|
||
x_proj = model[1](model[0](x)) # Linear + ReLU
|
||
gru_out, _ = model[2](x_proj) # GRU
|
||
output = model[3](gru_out) # Final Linear
|
||
else:
|
||
output = model(x)
|
||
|
||
# 同步TPU操作
|
||
xm.mark_step()
|
||
xm.wait_device_ops()
|
||
|
||
forward_time = time.time() - start_time
|
||
print(f"✅ 前向传播完成! 耗时: {forward_time:.3f}秒")
|
||
print(f"📤 输出形状: {output.shape}")
|
||
|
||
# 测试反向传播
|
||
print("🔄 测试反向传播...")
|
||
model.train()
|
||
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
|
||
|
||
start_time = time.time()
|
||
|
||
# 创建虚拟标签
|
||
labels = torch.randint(0, 41, (8, 50), device=device)
|
||
criterion = nn.CrossEntropyLoss()
|
||
|
||
# 前向传播
|
||
if hasattr(model, '__getitem__'):
|
||
x_proj = model[1](model[0](x))
|
||
gru_out, _ = model[2](x_proj)
|
||
output = model[3](gru_out)
|
||
else:
|
||
output = model(x)
|
||
|
||
# 计算损失
|
||
loss = criterion(output.view(-1, 41), labels.view(-1))
|
||
|
||
# 反向传播
|
||
optimizer.zero_grad()
|
||
loss.backward()
|
||
optimizer.step()
|
||
|
||
# 同步TPU操作
|
||
xm.mark_step()
|
||
xm.wait_device_ops()
|
||
|
||
backward_time = time.time() - start_time
|
||
print(f"✅ 反向传播完成! 耗时: {backward_time:.3f}秒")
|
||
print(f"🎯 损失值: {loss.item():.4f}")
|
||
|
||
# 总结
|
||
print(f"\n📈 性能总结:")
|
||
print(f" 前向传播: {forward_time:.3f}秒")
|
||
print(f" 反向传播: {backward_time:.3f}秒")
|
||
print(f" 总计: {forward_time + backward_time:.3f}秒")
|
||
|
||
if (forward_time + backward_time) < 10: # 10秒内完成
|
||
print("✅ TPU测试通过! 可以进行完整训练")
|
||
return True
|
||
else:
|
||
print("⚠️ TPU性能较慢，可能需要优化")
|
||
return False
|
||
|
||
except Exception as e:
|
||
print(f"❌ TPU测试失败: {e}")
|
||
import traceback
|
||
traceback.print_exc()
|
||
return False
|
||
|
||
|
||
if __name__ == "__main__":
|
||
print("=" * 50)
|
||
print("⚡ 快速TPU测试")
|
||
print("=" * 50)
|
||
|
||
success = quick_test()
|
||
|
||
if success:
|
||
print("\n🎉 测试成功! 现在可以运行:")
|
||
print(" python simple_tpu_model.py")
|
||
else:
|
||
print("\n❌ 测试失败，请检查TPU配置")
|
||
|
||
print("=" * 50) |