『PyTorch』使用指定 GPU 的方法

用户头像
kraken0
关注
发布于: 2020 年 06 月 02 日
『PyTorch』使用指定GPU的方法

PyTorch在不指定GPU序号时,默认使用id为0的GPU。



  1. 在终端中使用命令行直接设定:

CUDA_VISIBLE_DEVICES=id_of_GPU python your_file.py
  1. 在Python代码中指定:

import os
os.environ["CUDA_VISIBLE_DEVICES"] = "id_of_GPU"
  1. 使用PyTorch的函数torch.cuda.set_device:

import torch
torch.cuda.set_device(id)



查看当前使用GPU的序号:

import torch
torch.cuda.current_device()



用户头像

kraken0

关注

还未添加个人签名 2018.05.08 加入

还未添加个人简介

评论

发布
暂无评论
『PyTorch』使用指定GPU的方法