1 Pytorch中.to()和.cuda()的区别
如果需要指定的设备是GPU则.to()和.cuda()没有区别,如果设备是cpu,则不能使用.cuda()。也就是说.to()既可以指定CPU也可以指定GPU,而.cuda()只能指定GPU。
1.1 .cuda()
1.单GPU
<code class="language-python line-numbers">os.environ['<a href="https://www.stubbornhuang.com/tag/cuda/" title="浏览关于“CUDA”的文章" target="_blank" class="tag_link">CUDA</a>_VISIBLE_DEVICE']='0' model.cuda() </code>
2.多GPU
<code class="language-python line-numbers">os.environment['CUDA_VISIBLE_DEVICES'] = '0,1,2,3' device_ids = [0,1,2,3] model = torch.nn.Dataparallel(model, device_ids =device_ids) model = model.cuda() </code>
1.2 .to()
1.CPU或者单GPU
<code class="language-python line-numbers">device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") model.to(device) </code>
2.多GPU
<code class="language-python line-numbers">device_ids = [0,1,2,3] output_device = 0 model = nn.DataParallel(model,device_ids=device_ids) model.to(output_device) </code>
可以将上述代码进行封装
<code class="language-python line-numbers">class DataParallel(Module): def __init__(self, module, device_ids=None, output_device=None, dim=0): super(DataParallel, self).__init__() if not torch.cuda.is_available(): self.module = module self.device_ids = [] return if device_ids is None: device_ids = list(range(torch.cuda.device_count())) if output_device is None: output_device = device_ids[0] </code>
本文作者:StubbornHuang
版权声明:本文为站长原创文章,如果转载请注明原文链接!
原文标题:Pytorch – .to()和.cuda()的区别
原文链接:https://www.stubbornhuang.com/2202/
发布于:2022年07月11日 9:49:45
修改于:2023年06月25日 20:58:20
声明:本站所有文章,如无特殊说明或标注,均为本站原创发布。任何个人或组织,在未征得本站同意时,禁止复制、盗用、采集、发布本站内容到任何网站、书籍等各类媒体平台。如若本站内容侵犯了原著者的合法权益,可联系我们进行处理。
评论
49