File size: 2,503 Bytes
01c9370 | 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 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 | import argparse
import sys
import time
try:
import torch
except Exception as e:
print(f"导入PyTorch失败: {e}")
sys.exit(1)
def bytes_to_gb(x: int) -> float:
return x / 1024 ** 3
def clear_all_gpus(verbose: bool = True, sync: bool = True) -> None:
if not torch.cuda.is_available():
print("CUDA不可用，无需清理。")
return
num_gpus = torch.cuda.device_count()
if verbose:
print(f"检测到 {num_gpus} 个GPU，开始清理缓存…")
for i in range(num_gpus):
try:
torch.cuda.set_device(i)
before_alloc = torch.cuda.memory_allocated(i)
before_reserved = torch.cuda.memory_reserved(i)
# 清理缓存与IPC句柄
torch.cuda.empty_cache()
torch.cuda.ipc_collect()
if sync:
torch.cuda.synchronize(i)
after_alloc = torch.cuda.memory_allocated(i)
after_reserved = torch.cuda.memory_reserved(i)
if verbose:
print(
f"GPU {i}: 已分配 {bytes_to_gb(before_alloc):.2f}GB -> {bytes_to_gb(after_alloc):.2f}GB, "
f"已保留 {bytes_to_gb(before_reserved):.2f}GB -> {bytes_to_gb(after_reserved):.2f}GB"
)
except Exception as e:
print(f"清理GPU {i} 时出错: {e}")
if verbose:
print("清理完成。")
def main() -> int:
parser = argparse.ArgumentParser(description="清理所有GPU缓存与IPC资源")
parser.add_argument("--quiet", action="store_true", help="安静模式，不输出详细信息")
parser.add_argument("--no-sync", action="store_true", help="清理后不执行cuda同步")
parser.add_argument("--repeat", type=int, default=1, help="重复清理次数（默认1次）")
parser.add_argument(
"--interval", type=float, default=0.0, help="重复清理时每次间隔秒数（默认0）"
)
args = parser.parse_args()
verbose = not args.quiet
sync = not args.__dict__["no_sync"]
if not torch.cuda.is_available():
print("CUDA不可用，无需清理。")
return 0
for r in range(args.repeat):
if args.repeat > 1 and verbose:
print(f"第 {r + 1}/{args.repeat} 次清理…")
clear_all_gpus(verbose=verbose, sync=sync)
if r < args.repeat - 1 and args.interval > 0:
time.sleep(args.interval)
return 0
if __name__ == "__main__":
sys.exit(main())
|