【PyTorch】argparse + os.environ 设置pytorch网络使用的显卡
指定使用的显卡编号os.environ("CUDA_VISIBLE_DEVICES")='2,3,4'设置环境变量CUDA_VISIBLE_DEVICES为’2,3,4’,这个时候对于系统来说只有编号2,3,4的显卡是可见的(从0开始)通过torch.cuda.device_count()获取显卡数量的时候显示的是3,即只能看见这三张显卡在使用pytorch时,如果需要在gpu上对某些数据进行操作,一般的流程是:# 获取设备device = torch.device('cuda:0') # 使
原创
2022-05-26 15:47:13 ·
1717 阅读 ·
0 评论