torch.cuda.empty_cache()
 
# 設(shè)置進(jìn)程可使用的GPU顯存最大比例為50%
torch.cuda.set_per_process_memory_fraction(0.5, device=0)
 
# 計算總內(nèi)存
total_memory = torch.cuda.get_device_properties(0).total_memory
print("實際總內(nèi)存:", round(total_memory / (1024 * 1024), 1), "MB")
 
# 嘗試分配大量顯存的操作
try:
    # 使用10%的顯存:
    tmp_tensor = torch.empty(int(total_memory * 0.1), dtype=torch.int8, device='cuda:0')
    print("分配的內(nèi)存:", round(torch.cuda.memory_allocated(0) / (1024 * 1024), 1), "MB")
    print("保留的內(nèi)存:", round(torch.cuda.memory_reserved(0) / (1024 * 1024), 1), "MB")
    # 清空顯存
    del tmp_tensor
    torch.cuda.empty_cache()
    # 使用50%的顯存:
    torch.empty(int(total_memory * 0.5), dtype=torch.int8, device='cuda:0')
except RuntimeError as e:
    print("Error allocating tensor:", e)
 
# 打印當(dāng)前GPU的顯存使用情況
print("分配的內(nèi)存:", torch.cuda.memory_allocated(0) / (1024 * 1024), "MB")
print("保留的內(nèi)存:", torch.cuda.memory_reserved(0) / (1024 * 1024), "MB")

# 指定之后所有操作在 GPU3 上執(zhí)行
torch.cuda.set_device(3)

# 限制 GPU3 顯存使用50%
desired_memory_fraction = 0.5  # 50% 顯存
torch.cuda.set_per_process_memory_fraction(desired_memory_fraction)

# 獲取當(dāng)前GPU上的總顯存容量
total_memory = torch.cuda.get_device_properties(3).total_memory

# 指定使用 GPU3
tmp_tensor = torch.empty(int(total_memory * 0.4999), dtype=torch.int8, device="cuda") # 此處 cuda 即指 GPU3

# 獲取當(dāng)前已分配的顯存，計算可用顯存
allocated_memory = torch.cuda.memory_allocated()
available_memory = total_memory - allocated_memory

# 打印結(jié)果
print(f"Total GPU Memory: {total_memory / (1024**3):.2f} GB")
print(f"Allocated GPU Memory: {allocated_memory / (1024**3):.2f} GB")
print(f"Available GPU Memory: {available_memory / (1024**3):.2f} GB")

import torch
# 限制0號設(shè)備的顯存的使用量為0.5，就是半張卡那么多，比如12G卡，設(shè)置0.5就是6G。
torch.cuda.set_per_process_memory_fraction(0.5, 0)
torch.cuda.empty_cache()
# 計算一下總內(nèi)存有多少。
total_memory = torch.cuda.get_device_properties(0).total_memory
# 使用0.499的顯存:
tmp_tensor = torch.empty(int(total_memory * 0.499), dtype=torch.int8, device='cuda')

# 清空該顯存：
del tmp_tensor
torch.cuda.empty_cache()

# 下面這句話會觸發(fā)顯存OOM錯誤，因為剛好觸碰到了上限:
torch.empty(total_memory // 2, dtype=torch.int8, device='cuda')

"""
It raises an error as follows: 
RuntimeError: CUDA out of memory. Tried to allocate 5.59 GiB (GPU 0; 11.17 GiB total capacity; 0 bytes already allocated; 10.91 GiB free; 5.59 GiB allowed; 0 bytes reserved in total by PyTorch)
"""
顯存超標(biāo)后，比不設(shè)置限制的錯誤信息多了一個提示，“5.59 GiB allowed;”