PyTorch常用操作

指定GPU

  • 终端指定
1
CUDA_VISIBLE_DEVICES=0 nohup python demo.py >> base_log.out 2>&1 &
  • 程序指定
1
2
3
import torch
int id = {0-max_gpu_num}
torch.cuda.set_device(id)

检测GPU是否可用

1
2
import torch 
torch.cuda.is_available()