Pytorch開始作業前的前置檢查

檢查torch版本

import torch
torch.__version__

檢查cuda可用性/版本

import torch
torch.version.cuda #版本

torch.cuda.is_available() #有無

torch.cuda.device_count() #幾個device

torch.cuda.get_device_name() #型號

torch.cuda.get_device_properties(0).total_memory #記憶體大小
#這邊可以除(1024*1024*1024)，這樣看GB比較有感覺

torch.cuda.memory_allocated(0)/torch.cuda.get_device_properties(0).total_memory #目前被抓住的記憶大小

最重要的檢查(是否支援BF16):

torch.cuda.is_bf16_supported()

記得我們前面說過，A100之後才支援，這邊一樣可以驗證。