30 lines
961 B
Python
30 lines
961 B
Python
import torch
|
|
from loguru import logger
|
|
|
|
|
|
def peak_memory_decorator(func):
|
|
def wrapper(*args, **kwargs):
|
|
# 检查是否在分布式环境中
|
|
rank_info = ""
|
|
if torch.distributed.is_available() and torch.distributed.is_initialized():
|
|
rank = torch.distributed.get_rank()
|
|
rank_info = f"Rank {rank} - "
|
|
|
|
# 如果使用GPU，重置显存统计
|
|
if torch.cuda.is_available():
|
|
torch.cuda.reset_peak_memory_stats()
|
|
|
|
# 执行目标函数
|
|
result = func(*args, **kwargs)
|
|
|
|
# 获取峰值显存
|
|
if torch.cuda.is_available():
|
|
peak_memory = torch.cuda.max_memory_allocated() / (1024**3) # 转换为GB
|
|
logger.info(f"{rank_info}Function '{func.__qualname__}' Peak Memory: {peak_memory:.2f} GB")
|
|
else:
|
|
logger.info(f"{rank_info}Function '{func.__qualname__}' executed without GPU.")
|
|
|
|
return result
|
|
|
|
return wrapper
|