Files
ModelTrainEval/scripts/check_env.py
T
Shaun ZhouandClaude Opus 5 4b1147ed8f feat(scripts): 数据转换、ONNX 导出与推理性能分析
- convert_annotations: labelme / VOC -> YOLO txt,顺带生成 data.yaml
- split_dataset: 切分成 images/{train,val} 与 labels/{train,val},默认建软链不复制数据
- export_onnx: 导出部署用 ONNX,CPU 可跑,缺依赖时提示 `uv sync --group export`
- profile_inference: torch.profiler 测单帧延迟/FPS、CUDA 显存峰值与算子级明细
- check_env: Python / PyTorch / Nvidia GPU 自检

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
2026-09-02 17:55:20 +08:00

42 lines
1002 B
Python
Executable File

#!/usr/bin/env python
"""环境自检:确认依赖与 Nvidia GPU 是否就绪。
python scripts/check_env.py
"""
from __future__ import annotations
import sys
from pathlib import Path
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
from utils.device import NoNvidiaGpuError, cuda_devices, require_cuda # noqa: E402
def main() -> int:
print(f"Python: {sys.version.split()[0]}")
try:
import torch
print(f"PyTorch: {torch.__version__} (CUDA build: {torch.version.cuda})")
except ImportError:
print("PyTorch: 未安装 —— 先执行 `uv sync`")
return 1
devices = cuda_devices()
print(f"Nvidia GPU: {devices if devices else '无'}")
try:
require_cuda()
except NoNvidiaGpuError as exc:
print(f"训练不可用: {exc}")
print("(评估/导出仍可在 CPU 上进行)")
return 1
print("训练环境就绪。")
return 0
if __name__ == "__main__":
raise SystemExit(main())