File size: 929 Bytes
4811c23 | 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 | #!/usr/bin/env python
"""环境自检: GPU/CUDA/pyiqa/dwt/官方import 冒烟(不加载大权重)"""
import os, sys
sys.path.insert(0, "official"); sys.path.insert(0, "src")
import torch
print("torch", torch.__version__, "cuda", torch.cuda.is_available())
if torch.cuda.is_available():
print("gpu", torch.cuda.get_device_name(0), "mem_GB", round(torch.cuda.get_device_properties(0).total_memory/1e9,1))
try:
import pyiqa; print("pyiqa ok")
except Exception as e: print("pyiqa FAIL", e)
try:
import pywt; print("pywt ok")
except Exception as e: print("pywt missing(可选)", e)
try:
import diffusers, transformers, peft, omegaconf
print("diffusers/transformers/peft ok")
except Exception as e: print("libs FAIL", e)
# 官方模块可 import
try:
import model, dataset, utils, forward
print("official imports ok")
except Exception as e:
print("official imports FAIL", e)
print("env_check done")
|