当前位置:   article > 正文

pytorch数据选择CPU或者GPU_pytorch 1.13.1 py3.7_cpu_0 pytorch

pytorch 1.13.1 py3.7_cpu_0 pytorch

pytorch中.to(device) 和.cuda()的区别说明

原理

.to(device) 可以指定CPU 或者GPU

1

2

3

4

5

6

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") # 单GPU或者CPU

model.to(device)

#如果是多GPU

if torch.cuda.device_count() > 1:

  model = nn.DataParallel(model,device_ids=[0,1,2])

model.to(device)

.cuda() 只能指定GPU

1

2

3

4

5

6

7

8

9

#指定某个GPU

os.environ['CUDA_VISIBLE_DEVICE']='1'

model.cuda()

#如果是多GPU

os.environment['CUDA_VISIBLE_DEVICES'] = '0,1,2,3'

device_ids = [0,1,2,3]

net  = torch.nn.Dataparallel(net, device_ids =device_ids)

net  = torch.nn.Dataparallel(net) # 默认使用所有的device_ids

net = net.cuda()

采用.cuda的形式那么我们必须事先通过系统内部函数os.environment['CUDA_VISIBLE_DEVICES']

获取到系统的GPU

1

2

3

4

5

6

7

8

9

10

11

12

13

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]

补充:Pytorch使用To方法编写代码在不同设备(CUDA/CPU)上兼容(device-agnostic)

以前版本的PyTorch编写device-agnostic代码非常困难(即,在不修改代码的情况下在CUDA可以使用或者只能使用CPU的设备上运行)。

PyTorch 0.4.0使代码兼容

PyTorch 0.4.0通过两种方法使代码兼容变得非常容易:

张量的device属性为所有张量提供了torch.device设备。(注意:get_device仅适用于CUDA张量)

to方法Tensors和Modules可用于容易地将对象移动到不同的设备(代替以前的cpu()或cuda()方法)

我们推荐以下模式:

1

2

3

4

5

6

7

# 开始脚本,创建一个张量

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

...

# 但是无论你获得一个新的Tensor或者Module

# 如果他们已经在目标设备上则不会执行复制操作

input = data.to(device)

model = MyModule(...).to(device)

声明:本文内容由网友自发贡献,转载请注明出处:【wpsshop】
推荐阅读
相关标签
  

闽ICP备14008679号