Update main.py

main.py

@@ -9,47 +9,40 @@ from utils import load_image, load_mask, load_mask_edit
 from utils_mask import process_mask_move_torch, process_mask_remove_torch, mask_union_torch, mask_substract_torch, create_outer_edge_mask_torch
 from utils_mask import check_mask_overlap_torch, check_cover_all_torch, visualize_mask_list, get_mask_difference_torch, save_mask_list_to_npys
 
-print("PyTorch 版本:", torch.__version__)
-print("CUDA 版本:", torch.version.cuda)
-print("CUDA 是否可用:", torch.cuda.is_available())
-import subprocess
-
-# 检查 nvidia-smi 命令的输出
-try:
-    result = subprocess.run(['nvidia-smi'], stdout=subprocess.PIPE, stderr=subprocess.PIPE)
-    print(result.stdout.decode())
-except FileNotFoundError:
-    print("nvidia-smi 命令不可用，可能未安装 NVIDIA 驱动。")
-
-device_count = torch.cuda.device_count()
-print("可用的 CUDA 设备数量:", device_count)
-for i in range(device_count):
-    print(f"设备 {i} 名称:", torch.cuda.get_device_name(i))
-
 
 import os
-
-import glob
 import ctypes
+import torch
+
+# 设置 LD_LIBRARY_PATH
+nvidia_lib_path = '/usr/lib/x86_64-linux-gnu'
+os.environ['LD_LIBRARY_PATH'] = nvidia_lib_path + ':' + os.environ.get('LD_LIBRARY_PATH', '')
+print("LD_LIBRARY_PATH 已设置为:", os.environ['LD_LIBRARY_PATH'])
+
+# 尝试加载 libnvidia-ml.so
 try:
     ctypes.CDLL('libnvidia-ml.so')
     print("成功加载 libnvidia-ml.so。")
 except Exception as e:
     print("无法加载 libnvidia-ml.so，错误信息:", e)
 
-libs = glob.glob('/usr/lib*/**/libnvidia-*.so*', recursive=True)
-print("NVIDIA 驱动库文件:", libs)
-print("LD_LIBRARY_PATH:", os.environ.get('LD_LIBRARY_PATH', '未设置'))
+# 尝试加载 libcuda.so
+try:
+    ctypes.CDLL('libcuda.so')
+    print("成功加载 libcuda.so。")
+except Exception as e:
+    print("无法加载 libcuda.so，错误信息:", e)
+
+# 检查 CUDA 是否可用
+print("CUDA 是否可用:", torch.cuda.is_available())
+print("可用的 CUDA 设备数量:", torch.cuda.device_count())
+
+# 尝试在 CUDA 上创建张量
 try:
-    with open('/proc/driver/nvidia/version') as f:
-        print(f.read())
+    x = torch.tensor([1.0], device='cuda')
+    print("CUDA 张量创建成功:", x)
 except Exception as e:
-    print("无法读取驱动程序版本信息，错误信息:", e)
-print("PyTorch 安装路径:", torch.__file__)
-cuda_paths = [path for path in os.listdir('/usr/local') if path.startswith('cuda')]
-print("CUDA 安装路径:", cuda_paths)
-print("PyTorch CUDA 版本:", torch.version.cuda)
-print("系统 CUDA 版本:", os.environ.get('CUDA_VERSION', '未设置'))
+    print("无法在 CUDA 上创建张量，错误信息:", e)
 
     
 def run_main(