663 lines
23 KiB
Python
663 lines
23 KiB
Python
#!/usr/bin/env python3
|
||
"""
|
||
性能基准测试脚本
|
||
用于诊断输入法模型训练的性能瓶颈
|
||
"""
|
||
|
||
import json
|
||
import subprocess
|
||
import sys
|
||
import time
|
||
from pathlib import Path
|
||
from typing import Any, Dict, List, Tuple
|
||
|
||
import torch
|
||
import torch.nn as nn
|
||
import torch.nn.functional as F
|
||
from torch.utils.data import DataLoader, Dataset
|
||
|
||
# 添加src目录到路径
|
||
sys.path.insert(0, str(Path(__file__).parent))
|
||
|
||
from src.model.model import InputMethodEngine
|
||
|
||
|
||
class MockDataset(Dataset):
|
||
"""模拟数据集用于性能测试"""
|
||
|
||
def __init__(self, num_samples=1000, vocab_size=10019, seq_len=128, num_slots=8):
|
||
self.num_samples = num_samples
|
||
self.vocab_size = vocab_size
|
||
self.seq_len = seq_len
|
||
self.num_slots = num_slots
|
||
|
||
def __len__(self):
|
||
return self.num_samples
|
||
|
||
def __getitem__(self, idx):
|
||
# 生成模拟数据，包含训练所需的所有字段
|
||
return {
|
||
"input_ids": torch.randint(0, self.vocab_size, (self.seq_len,)),
|
||
"token_type_ids": torch.randint(0, 2, (self.seq_len,)),
|
||
"attention_mask": torch.ones(self.seq_len, dtype=torch.long),
|
||
"history_slot_ids": torch.randint(0, self.vocab_size, (self.num_slots,)),
|
||
"pinyin_ids": torch.randint(0, 30, (24,)), # pinyin_ids长度固定为24
|
||
"label": torch.randint(0, self.vocab_size, (1,)),
|
||
}
|
||
|
||
|
||
def collate_fn(batch: List[Dict[str, Any]]) -> Dict[str, Any]:
|
||
"""与trainer.py中相同的collate_fn"""
|
||
# 处理tensor字段
|
||
input_ids = torch.stack([item["input_ids"] for item in batch])
|
||
token_type_ids = torch.stack([item["token_type_ids"] for item in batch])
|
||
attention_mask = torch.stack([item["attention_mask"] for item in batch])
|
||
labels = torch.stack([item["label"] for item in batch])
|
||
history_slot_ids = torch.stack([item["history_slot_ids"] for item in batch])
|
||
pinyin_ids = torch.stack([item["pinyin_ids"] for item in batch])
|
||
|
||
# 字符串字段（模拟）
|
||
prefixes = [f"prefix_{i}" for i in range(len(batch))]
|
||
suffixes = [f"suffix_{i}" for i in range(len(batch))]
|
||
pinyins = [f"pinyin_{i}" for i in range(len(batch))]
|
||
|
||
return {
|
||
"input_ids": input_ids,
|
||
"token_type_ids": token_type_ids,
|
||
"attention_mask": attention_mask,
|
||
"labels": labels,
|
||
"history_slot_ids": history_slot_ids,
|
||
"pinyin_ids": pinyin_ids,
|
||
"prefix": prefixes,
|
||
"suffix": suffixes,
|
||
"pinyin": pinyins,
|
||
}
|
||
|
||
|
||
def get_gpu_utilization() -> float:
|
||
"""获取GPU利用率"""
|
||
try:
|
||
result = subprocess.run(
|
||
[
|
||
"nvidia-smi",
|
||
"--query-gpu=utilization.gpu",
|
||
"--format=csv,noheader,nounits",
|
||
],
|
||
capture_output=True,
|
||
text=True,
|
||
timeout=2,
|
||
)
|
||
if result.returncode == 0:
|
||
return float(result.stdout.strip())
|
||
except:
|
||
pass
|
||
return 0.0
|
||
|
||
|
||
def benchmark_model_only(batch_sizes=[32, 64, 128, 256, 512]):
|
||
"""仅测试模型前向+反向传播性能（不包含数据加载）"""
|
||
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
||
print(f"\n{'=' * 70}")
|
||
print("模型计算性能测试（不包含数据加载）")
|
||
print(f"使用设备: {device}")
|
||
print(f"PyTorch版本: {torch.__version__}")
|
||
print(f"CUDA可用: {torch.cuda.is_available()}")
|
||
|
||
if torch.cuda.is_available():
|
||
print(f"GPU: {torch.cuda.get_device_name(0)}")
|
||
print(
|
||
f"GPU内存: {torch.cuda.get_device_properties(0).total_memory / 1024**3:.2f} GB"
|
||
)
|
||
|
||
# 关键修复：使用与ContextEncoder预训练模型匹配的维度
|
||
# 根据components.py，ContextEncoder使用预训练模型，其embedding维度是固定的
|
||
# 从错误信息看，预训练模型维度是512，所以这里必须使用dim=512
|
||
vocab_size = 10019
|
||
pinyin_vocab_size = 30
|
||
dim = 512 # 必须与预训练模型维度匹配！
|
||
num_slots = 8
|
||
n_layers = 4
|
||
n_heads = 4
|
||
num_experts = 20
|
||
max_seq_len = 128
|
||
|
||
results = []
|
||
|
||
for batch_size in batch_sizes:
|
||
print(f"\n{'=' * 60}")
|
||
print(f"测试批量大小: {batch_size}")
|
||
|
||
try:
|
||
# 创建模型（使用真实训练配置）
|
||
model = InputMethodEngine(
|
||
vocab_size=vocab_size,
|
||
pinyin_vocab_size=pinyin_vocab_size,
|
||
dim=dim,
|
||
num_slots=num_slots,
|
||
n_layers=n_layers,
|
||
n_heads=n_heads,
|
||
num_experts=num_experts,
|
||
max_seq_len=max_seq_len,
|
||
).to(device)
|
||
|
||
# 创建优化器和损失函数
|
||
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4)
|
||
criterion = nn.CrossEntropyLoss()
|
||
scaler = torch.cuda.amp.GradScaler(enabled=torch.cuda.is_available())
|
||
|
||
# 生成模拟batch（直接在GPU上）
|
||
batch = {
|
||
"input_ids": torch.randint(
|
||
0, vocab_size, (batch_size, max_seq_len), device=device
|
||
),
|
||
"token_type_ids": torch.randint(
|
||
0, 2, (batch_size, max_seq_len), device=device
|
||
),
|
||
"attention_mask": torch.ones(
|
||
batch_size, max_seq_len, dtype=torch.long, device=device
|
||
),
|
||
"history_slot_ids": torch.randint(
|
||
0, vocab_size, (batch_size, num_slots), device=device
|
||
),
|
||
"pinyin_ids": torch.randint(
|
||
0, pinyin_vocab_size, (batch_size, 24), device=device
|
||
),
|
||
"labels": torch.randint(0, vocab_size, (batch_size, 1), device=device),
|
||
}
|
||
|
||
# 预热（运行几次以避免初始化开销）
|
||
print(" 预热...")
|
||
for _ in range(3):
|
||
with torch.cuda.amp.autocast(enabled=torch.cuda.is_available()):
|
||
logits = model(
|
||
input_ids=batch["input_ids"],
|
||
token_type_ids=batch["token_type_ids"],
|
||
attention_mask=batch["attention_mask"],
|
||
pinyin_ids=batch["pinyin_ids"],
|
||
history_slot_ids=batch["history_slot_ids"],
|
||
)
|
||
loss = criterion(logits, batch["labels"].squeeze(-1))
|
||
scaler.scale(loss).backward()
|
||
scaler.step(optimizer)
|
||
scaler.update()
|
||
optimizer.zero_grad()
|
||
|
||
# 测速
|
||
if torch.cuda.is_available():
|
||
torch.cuda.synchronize()
|
||
start_time = time.time()
|
||
steps = 20 # 运行20步以获得稳定平均值
|
||
|
||
print(f" 运行{steps}步测速...")
|
||
for step in range(steps):
|
||
with torch.cuda.amp.autocast(enabled=torch.cuda.is_available()):
|
||
logits = model(
|
||
input_ids=batch["input_ids"],
|
||
token_type_ids=batch["token_type_ids"],
|
||
attention_mask=batch["attention_mask"],
|
||
pinyin_ids=batch["pinyin_ids"],
|
||
history_slot_ids=batch["history_slot_ids"],
|
||
)
|
||
loss = criterion(logits, batch["labels"].squeeze(-1))
|
||
scaler.scale(loss).backward()
|
||
scaler.step(optimizer)
|
||
scaler.update()
|
||
optimizer.zero_grad()
|
||
|
||
if torch.cuda.is_available():
|
||
torch.cuda.synchronize()
|
||
elapsed = time.time() - start_time
|
||
|
||
# 计算指标
|
||
throughput = steps * batch_size / elapsed
|
||
step_time = elapsed / steps * 1000 # 毫秒
|
||
memory_used = (
|
||
torch.cuda.max_memory_allocated() / 1024**2
|
||
if torch.cuda.is_available()
|
||
else 0
|
||
) # MB
|
||
|
||
print(f" 结果:")
|
||
print(f" 吞吐量: {throughput:.2f} samples/sec")
|
||
print(f" 每个step: {step_time:.2f}ms")
|
||
print(f" GPU内存峰值: {memory_used:.2f} MB")
|
||
|
||
if torch.cuda.is_available():
|
||
gpu_util = get_gpu_utilization()
|
||
if gpu_util > 0:
|
||
print(f" GPU利用率: {gpu_util:.1f}%")
|
||
|
||
results.append(
|
||
{
|
||
"batch_size": batch_size,
|
||
"throughput": throughput,
|
||
"step_time": step_time,
|
||
"memory": memory_used,
|
||
"success": True,
|
||
}
|
||
)
|
||
|
||
# 清理
|
||
del model, optimizer, batch
|
||
if torch.cuda.is_available():
|
||
torch.cuda.empty_cache()
|
||
|
||
# 检查是否内存不足
|
||
if (
|
||
torch.cuda.is_available()
|
||
and memory_used
|
||
> torch.cuda.get_device_properties(0).total_memory * 0.8 / 1024**2
|
||
):
|
||
print(
|
||
f" ⚠️ 警告: batch_size={batch_size} 使用了超过80% GPU内存，更大的batch_size可能导致OOM"
|
||
)
|
||
break
|
||
|
||
except Exception as e:
|
||
print(f" 错误: {e}")
|
||
results.append(
|
||
{
|
||
"batch_size": batch_size,
|
||
"error": str(e),
|
||
"success": False,
|
||
}
|
||
)
|
||
if "CUDA out of memory" in str(e):
|
||
print(f" ⚠️ CUDA内存不足，停止测试更大的batch_size")
|
||
break
|
||
|
||
return results
|
||
|
||
|
||
def benchmark_data_loading(num_workers_list=[0, 2, 4, 8, 12], batch_size=128):
|
||
"""测试数据加载性能"""
|
||
print(f"\n{'=' * 70}")
|
||
print("测试数据加载性能")
|
||
print(f"批量大小: {batch_size}")
|
||
|
||
results = []
|
||
|
||
for num_workers in num_workers_list:
|
||
print(f"\n测试num_workers: {num_workers}")
|
||
|
||
try:
|
||
# 创建数据集和数据加载器
|
||
dataset = MockDataset(num_samples=2000)
|
||
dataloader = DataLoader(
|
||
dataset,
|
||
batch_size=batch_size,
|
||
num_workers=num_workers,
|
||
pin_memory=torch.cuda.is_available(),
|
||
collate_fn=collate_fn,
|
||
prefetch_factor=2 if num_workers > 0 else None,
|
||
)
|
||
|
||
# 预热
|
||
print(" 预热...")
|
||
for _ in range(2):
|
||
_ = next(iter(dataloader))
|
||
|
||
# 测速
|
||
print(f" 测速中...")
|
||
start_time = time.time()
|
||
batches = 0
|
||
total_samples = 0
|
||
|
||
for batch in dataloader:
|
||
batches += 1
|
||
total_samples += batch["input_ids"].size(0)
|
||
|
||
if batches >= 20: # 测试20个batch
|
||
break
|
||
|
||
elapsed = time.time() - start_time
|
||
if batches == 0:
|
||
print(f" 警告: 无法加载任何batch")
|
||
continue
|
||
|
||
throughput = total_samples / elapsed
|
||
|
||
print(f" 结果:")
|
||
print(f" 数据加载吞吐量: {throughput:.2f} samples/sec")
|
||
print(f" 平均每batch加载时间: {elapsed / batches * 1000:.2f}ms")
|
||
|
||
results.append(
|
||
{
|
||
"num_workers": num_workers,
|
||
"throughput": throughput,
|
||
"avg_batch_time": elapsed / batches * 1000,
|
||
"success": True,
|
||
}
|
||
)
|
||
|
||
except Exception as e:
|
||
print(f" 错误: {e}")
|
||
results.append(
|
||
{
|
||
"num_workers": num_workers,
|
||
"error": str(e),
|
||
"success": False,
|
||
}
|
||
)
|
||
|
||
return results
|
||
|
||
|
||
def profile_training_step(batch_size=128, num_steps=15):
|
||
"""详细分析训练步骤中各部分耗时"""
|
||
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
||
|
||
print(f"\n{'=' * 70}")
|
||
print("详细训练步骤分析")
|
||
print(f"批量大小: {batch_size}, 步数: {num_steps}")
|
||
|
||
# 使用真实训练配置
|
||
vocab_size = 10019
|
||
pinyin_vocab_size = 30
|
||
dim = 512
|
||
num_slots = 8
|
||
n_layers = 4
|
||
n_heads = 4
|
||
num_experts = 20
|
||
max_seq_len = 128
|
||
|
||
timings = {
|
||
"data_to_gpu": [],
|
||
"forward": [],
|
||
"loss_calc": [],
|
||
"backward": [],
|
||
"optimizer_step": [],
|
||
"total_step": [],
|
||
}
|
||
|
||
try:
|
||
# 创建模型和优化器
|
||
model = InputMethodEngine(
|
||
vocab_size=vocab_size,
|
||
pinyin_vocab_size=pinyin_vocab_size,
|
||
dim=dim,
|
||
num_slots=num_slots,
|
||
n_layers=n_layers,
|
||
n_heads=n_heads,
|
||
num_experts=num_experts,
|
||
max_seq_len=max_seq_len,
|
||
).to(device)
|
||
|
||
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4)
|
||
criterion = nn.CrossEntropyLoss()
|
||
scaler = torch.cuda.amp.GradScaler(enabled=torch.cuda.is_available())
|
||
|
||
# 创建数据加载器
|
||
dataset = MockDataset(num_samples=1000)
|
||
dataloader = DataLoader(
|
||
dataset,
|
||
batch_size=batch_size,
|
||
num_workers=4,
|
||
pin_memory=torch.cuda.is_available(),
|
||
collate_fn=collate_fn,
|
||
)
|
||
|
||
model.train()
|
||
|
||
print(f"开始分析...")
|
||
for i, batch in enumerate(dataloader):
|
||
if i >= num_steps:
|
||
break
|
||
|
||
step_start = time.time()
|
||
|
||
# 1. 数据移动到GPU
|
||
data_start = time.time()
|
||
batch_gpu = {
|
||
k: v.to(device, non_blocking=True)
|
||
for k, v in batch.items()
|
||
if isinstance(v, torch.Tensor)
|
||
}
|
||
data_time = time.time() - data_start
|
||
timings["data_to_gpu"].append(data_time)
|
||
|
||
# 2. 前向传播
|
||
forward_start = time.time()
|
||
with torch.cuda.amp.autocast(enabled=torch.cuda.is_available()):
|
||
logits = model(
|
||
input_ids=batch_gpu["input_ids"],
|
||
token_type_ids=batch_gpu["token_type_ids"],
|
||
attention_mask=batch_gpu["attention_mask"],
|
||
pinyin_ids=batch_gpu["pinyin_ids"],
|
||
history_slot_ids=batch_gpu["history_slot_ids"],
|
||
)
|
||
forward_time = time.time() - forward_start
|
||
timings["forward"].append(forward_time)
|
||
|
||
# 3. 损失计算
|
||
loss_start = time.time()
|
||
with torch.cuda.amp.autocast(enabled=torch.cuda.is_available()):
|
||
loss = criterion(logits, batch_gpu["labels"].squeeze(-1))
|
||
loss_time = time.time() - loss_start
|
||
timings["loss_calc"].append(loss_time)
|
||
|
||
# 4. 反向传播
|
||
backward_start = time.time()
|
||
scaler.scale(loss).backward()
|
||
backward_time = time.time() - backward_start
|
||
timings["backward"].append(backward_time)
|
||
|
||
# 5. 优化器步骤
|
||
optimizer_start = time.time()
|
||
scaler.step(optimizer)
|
||
scaler.update()
|
||
optimizer.zero_grad()
|
||
optimizer_time = time.time() - optimizer_start
|
||
timings["optimizer_step"].append(optimizer_time)
|
||
|
||
# 总时间
|
||
total_time = time.time() - step_start
|
||
timings["total_step"].append(total_time)
|
||
|
||
if i % 5 == 0:
|
||
print(
|
||
f" Step {i}: 总时间={total_time * 1000:.1f}ms, "
|
||
f"前向={forward_time * 1000:.1f}ms, "
|
||
f"反向={backward_time * 1000:.1f}ms"
|
||
)
|
||
|
||
# 打印统计信息
|
||
print(f"\n{'=' * 60}")
|
||
print("训练步骤耗时分析（平均值）：")
|
||
for key, values in timings.items():
|
||
if values:
|
||
avg_ms = sum(values) / len(values) * 1000
|
||
min_ms = min(values) * 1000
|
||
max_ms = max(values) * 1000
|
||
if timings["total_step"]:
|
||
total_avg = (
|
||
sum(timings["total_step"]) / len(timings["total_step"]) * 1000
|
||
)
|
||
percentage = (avg_ms / total_avg) * 100
|
||
else:
|
||
percentage = 0
|
||
print(
|
||
f" {key:15s}: {avg_ms:6.1f}ms ({percentage:5.1f}%) [min={min_ms:.1f}ms, max={max_ms:.1f}ms]"
|
||
)
|
||
|
||
# 计算瓶颈
|
||
if timings["total_step"]:
|
||
total_avg_time = sum(timings["total_step"]) / len(timings["total_step"])
|
||
data_percent = (
|
||
sum(timings["data_to_gpu"]) / sum(timings["total_step"]) * 100
|
||
)
|
||
compute_percent = (
|
||
(
|
||
sum(timings["forward"])
|
||
+ sum(timings["backward"])
|
||
+ sum(timings["loss_calc"])
|
||
+ sum(timings["optimizer_step"])
|
||
)
|
||
/ sum(timings["total_step"])
|
||
* 100
|
||
)
|
||
|
||
print(f"\n瓶颈分析：")
|
||
print(f" 数据加载/传输: {data_percent:.1f}%")
|
||
print(f" 计算（前向+反向）: {compute_percent:.1f}%")
|
||
|
||
if data_percent > 30:
|
||
print(f" ⚠️ 数据加载是瓶颈！建议增加num_workers或使用pin_memory")
|
||
if compute_percent > 70:
|
||
print(f" ⚠️ 计算是瓶颈！建议优化模型或使用混合精度")
|
||
|
||
except Exception as e:
|
||
print(f"分析过程中出错: {e}")
|
||
import traceback
|
||
|
||
traceback.print_exc()
|
||
|
||
return timings
|
||
|
||
|
||
def check_gpu_status():
|
||
"""检查GPU状态"""
|
||
print(f"\n{'=' * 70}")
|
||
print("GPU状态检查")
|
||
|
||
if not torch.cuda.is_available():
|
||
print("CUDA不可用！将在CPU上运行")
|
||
print("注意：CPU训练性能会远低于GPU")
|
||
return
|
||
|
||
print(f"可用GPU数量: {torch.cuda.device_count()}")
|
||
|
||
for i in range(torch.cuda.device_count()):
|
||
props = torch.cuda.get_device_properties(i)
|
||
print(f"\nGPU {i}: {props.name}")
|
||
print(f" 计算能力: {props.major}.{props.minor}")
|
||
print(f" 总内存: {props.total_memory / 1024**3:.2f} GB")
|
||
print(f" SM数量: {props.multi_processor_count}")
|
||
|
||
# 当前内存使用
|
||
print(f"\n当前内存使用:")
|
||
print(f" 已分配: {torch.cuda.memory_allocated() / 1024**2:.2f} MB")
|
||
print(f" 已缓存: {torch.cuda.memory_reserved() / 1024**2:.2f} MB")
|
||
|
||
# 检查PyTorch配置
|
||
print(f"\nPyTorch配置:")
|
||
print(f" cuDNN基准模式: {torch.backends.cudnn.benchmark}")
|
||
print(f" cuDNN确定性: {torch.backends.cudnn.deterministic}")
|
||
print(f" TF32启用: {torch.backends.cuda.matmul.allow_tf32}")
|
||
|
||
# 建议优化
|
||
print(f"\n建议优化:")
|
||
if not torch.backends.cudnn.benchmark:
|
||
print(f" ⚙️ 建议设置: torch.backends.cudnn.benchmark = True")
|
||
if not torch.backends.cuda.matmul.allow_tf32:
|
||
print(f" ⚙️ 建议设置: torch.backends.cuda.matmul.allow_tf32 = True")
|
||
|
||
|
||
def main():
|
||
"""主函数：运行所有性能测试"""
|
||
print("输入法模型性能基准测试")
|
||
print("=" * 70)
|
||
|
||
# 设置优化标志（在测试前设置）
|
||
if torch.cuda.is_available():
|
||
torch.backends.cudnn.benchmark = True
|
||
torch.backends.cuda.matmul.allow_tf32 = True
|
||
print("已启用cuDNN基准模式和TF32加速")
|
||
|
||
# 1. 检查GPU状态
|
||
check_gpu_status()
|
||
|
||
# 2. 测试数据加载性能
|
||
print(f"\n{'=' * 70}")
|
||
print("阶段1: 数据加载性能测试")
|
||
data_results = benchmark_data_loading(
|
||
num_workers_list=[0, 2, 4, 8, 12], batch_size=128
|
||
)
|
||
|
||
# 3. 测试模型计算性能
|
||
print(f"\n{'=' * 70}")
|
||
print("阶段2: 模型计算性能测试")
|
||
# 从较小的batch_size开始测试
|
||
model_results = benchmark_model_only(batch_sizes=[32, 64, 128, 256])
|
||
|
||
# 4. 详细分析训练步骤
|
||
print(f"\n{'=' * 70}")
|
||
print("阶段3: 详细训练步骤分析")
|
||
step_timings = profile_training_step(batch_size=128, num_steps=15)
|
||
|
||
# 5. 给出优化建议
|
||
print(f"\n{'=' * 70}")
|
||
print("优化建议汇总")
|
||
print("=" * 70)
|
||
|
||
# 基于数据加载结果
|
||
if data_results:
|
||
successful_results = [r for r in data_results if r.get("success", False)]
|
||
if successful_results:
|
||
best_workers = max(successful_results, key=lambda x: x["throughput"])
|
||
print(f"1. 数据加载优化:")
|
||
print(f" 推荐num_workers: {best_workers['num_workers']}")
|
||
print(f" 最佳吞吐量: {best_workers['throughput']:.2f} samples/sec")
|
||
|
||
# 基于模型结果
|
||
if model_results:
|
||
successful_results = [r for r in model_results if r.get("success", False)]
|
||
if successful_results:
|
||
best_batch = max(successful_results, key=lambda x: x["throughput"])
|
||
print(f"\n2. 批量大小优化:")
|
||
print(f" 推荐batch_size: {best_batch['batch_size']}")
|
||
print(f" 最佳吞吐量: {best_batch['throughput']:.2f} samples/sec")
|
||
print(f" 内存使用: {best_batch['memory']:.2f} MB")
|
||
|
||
# 内存安全建议
|
||
if best_batch["memory"] > 0:
|
||
if torch.cuda.is_available():
|
||
total_mem = (
|
||
torch.cuda.get_device_properties(0).total_memory / 1024**2
|
||
)
|
||
usage_percent = (best_batch["memory"] / total_mem) * 100
|
||
if usage_percent < 70:
|
||
print(f" ✅ 内存使用安全 ({usage_percent:.1f}% of total)")
|
||
else:
|
||
print(
|
||
f" ⚠️ 内存使用较高 ({usage_percent:.1f}% of total)，考虑减小batch_size"
|
||
)
|
||
|
||
# 通用建议
|
||
print(f"\n3. 通用优化建议:")
|
||
if torch.cuda.is_available():
|
||
print(f" ✅ 使用pin_memory=True (已启用)")
|
||
print(f" ✅ 使用混合精度训练 (已启用)")
|
||
print(f" ✅ 设置torch.backends.cudnn.benchmark = True (已启用)")
|
||
print(f" ✅ 使用非阻塞数据传输 (non_blocking=True)")
|
||
else:
|
||
print(f" ⚠️ 未检测到GPU，建议使用GPU以获得更好性能")
|
||
|
||
# 特定建议
|
||
print(f"\n4. 针对你的配置的具体建议:")
|
||
print(f" - 默认配置: batch_size=128, num_workers=2, dim=512")
|
||
print(f" - 如果GPU内存充足，可以尝试增加batch_size")
|
||
print(f" - 如果CPU核心多，可以增加num_workers")
|
||
|
||
print(f"\n{'=' * 70}")
|
||
print("基准测试完成！")
|
||
print(f"可以运行完整训练来验证优化效果:")
|
||
print(f" python -m src.model.trainer train --help")
|
||
|
||
# 保存结果到文件
|
||
try:
|
||
results = {
|
||
"data_loading": data_results,
|
||
"model_compute": model_results,
|
||
"timestamp": time.strftime("%Y-%m-%d %H:%M:%S"),
|
||
}
|
||
with open("benchmark_results.json", "w") as f:
|
||
json.dump(results, f, indent=2, default=str)
|
||
print(f"\n结果已保存到: benchmark_results.json")
|
||
except:
|
||
pass
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|