PyTorch中如何使用指定的GPU

Song712 次浏览0个评论2018年07月11日
描述

PyTorch默认使用从0开始的GPU,如果GPU0正在运行程序,需要指定其他GPU。

[attach]26[/attach]

有如下两种方法来指定需要使用的GPU。 [list] []使用[b]CUDA_VISIBLE_DEVICES[/b],类似tensorflow指定GPU的方法。[/] []直接终端中设定:[code]CUDA_VISIBLE_DEVICES=1 python my_script.py[/code][/] []python代码中设定:[code]import os os.environ["CUDA_VISIBLE_DEVICES"] = "2"[/code][/] []使用函数 [b]set_device,[/b]该函数见 pytorch-master\torch\cuda__init__.py。[code]import torch torch.cuda.set_device(id)[/code][/] [/list]   不过官方建议使用[b]CUDA_VISIBLE_DEVICES[/b],不建议使用 [b]set_device[/b] 函数。

提交评论

请登录后评论

用户评论

    当前暂无评价,快来发表您的观点吧...

更多相关好文

    当前暂无更多相关好文推荐...

    公告提示

    Pytorch是Facebook的AI研究团队发布了一个Python工具包,是Python优先的深度学习框架。作为numpy的替代品;使用强大的GPU能力,提供最大的灵活性和速度,实现了机器学习框架Torch在Python语言环境的执行,基于python且具备强大GPU加速的张量和动态神经网络。