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

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)

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]

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)