124 lines
3.8 KiB
Python
124 lines
3.8 KiB
Python
# ====================
|
||
# 单元格4: 逐步调试完整模型编译
|
||
# ====================
|
||
|
||
# 如果单元格3测试通过，运行这个单元格
|
||
print("🔧 逐步测试完整TripleGRUDecoder模型...")
|
||
|
||
# 导入完整模型
|
||
import sys
|
||
sys.path.append('.') # 确保能导入本地模块
|
||
|
||
try:
|
||
from rnn_model import TripleGRUDecoder
|
||
print("✅ TripleGRUDecoder导入成功")
|
||
except ImportError as e:
|
||
print(f"❌ 模型导入失败: {e}")
|
||
print("请确保rnn_model.py在当前目录中")
|
||
|
||
# 分阶段测试
|
||
def test_model_compilation_stages():
|
||
"""分阶段测试模型编译"""
|
||
device = xm.xla_device()
|
||
|
||
# 阶段1: 测试NoiseModel单独编译
|
||
print("\n🔬 阶段1: 测试NoiseModel...")
|
||
try:
|
||
from rnn_model import NoiseModel
|
||
noise_model = NoiseModel(
|
||
neural_dim=512,
|
||
n_units=384, # 减小参数
|
||
n_days=4,
|
||
patch_size=8 # 减小patch size
|
||
).to(device)
|
||
|
||
x = torch.randn(2, 20, 512, device=device)
|
||
day_idx = torch.tensor([0, 1], device=device)
|
||
|
||
start_time = time.time()
|
||
with torch.no_grad():
|
||
output, states = noise_model(x, day_idx)
|
||
compile_time = time.time() - start_time
|
||
|
||
print(f"✅ NoiseModel编译成功! 耗时: {compile_time:.2f}秒")
|
||
print(f" 参数数量: {sum(p.numel() for p in noise_model.parameters()):,}")
|
||
|
||
return True, compile_time
|
||
|
||
except Exception as e:
|
||
print(f"❌ NoiseModel编译失败: {e}")
|
||
return False, 0
|
||
|
||
# 阶段2: 测试CleanSpeechModel
|
||
print("\n🔬 阶段2: 测试CleanSpeechModel...")
|
||
try:
|
||
from rnn_model import CleanSpeechModel
|
||
clean_model = CleanSpeechModel(
|
||
neural_dim=512,
|
||
n_units=384,
|
||
n_days=4,
|
||
n_classes=41,
|
||
patch_size=8
|
||
).to(device)
|
||
|
||
start_time = time.time()
|
||
with torch.no_grad():
|
||
output = clean_model(x, day_idx)
|
||
compile_time = time.time() - start_time
|
||
|
||
print(f"✅ CleanSpeechModel编译成功! 耗时: {compile_time:.2f}秒")
|
||
return True, compile_time
|
||
|
||
except Exception as e:
|
||
print(f"❌ CleanSpeechModel编译失败: {e}")
|
||
return False, 0
|
||
|
||
# 阶段3: 测试完整TripleGRUDecoder
|
||
print("\n🔬 阶段3: 测试TripleGRUDecoder...")
|
||
try:
|
||
model = TripleGRUDecoder(
|
||
neural_dim=512,
|
||
n_units=384, # 比原来的768小
|
||
n_days=4, # 减少天数
|
||
n_classes=41,
|
||
patch_size=8 # 减小patch size
|
||
).to(device)
|
||
|
||
print(f"📊 完整模型参数: {sum(p.numel() for p in model.parameters()):,}")
|
||
|
||
# 启动编译监控
|
||
compilation_monitor.start_monitoring()
|
||
|
||
start_time = time.time()
|
||
with torch.no_grad():
|
||
# 测试inference模式
|
||
logits = model(x, day_idx, None, False, 'inference')
|
||
compile_time = time.time() - start_time
|
||
|
||
compilation_monitor.complete_monitoring()
|
||
|
||
print(f"✅ TripleGRUDecoder编译成功! 耗时: {compile_time:.2f}秒")
|
||
print(f"📤 输出形状: {logits.shape}")
|
||
|
||
return True, compile_time
|
||
|
||
except Exception as e:
|
||
compilation_monitor.complete_monitoring()
|
||
print(f"❌ TripleGRUDecoder编译失败: {e}")
|
||
return False, 0
|
||
|
||
# 运行分阶段测试
|
||
stage_results = test_model_compilation_stages()
|
||
|
||
if stage_results:
|
||
print(f"\n🎉 所有编译测试完成!")
|
||
print("💡 下一步可以尝试:")
|
||
print(" 1. 使用简化配置进行训练")
|
||
print(" 2. 逐步增加模型复杂度")
|
||
print(" 3. 监控TPU资源使用情况")
|
||
else:
|
||
print(f"\n⚠️ 编译测试发现问题")
|
||
print("💡 建议:")
|
||
print(" 1. 进一步减小模型参数")
|
||
print(" 2. 检查内存使用情况")
|
||
print(" 3. 使用CPU模式进行调试") |