Tensorflow, pyTorch - 원하는 GPU 할당

모델 학습을 하기 위해 원하는 GPU만 사용하기 위해서는

python의 os 모듈을 통해 간단히 접근할 수 있는 GPU 번호를 제한해주면 된다.

 

 

1. 1개의 단일 GPU 사용

import os
os.environ['CUDA_VISIBLE_DEVICES'] = '0'

 

2. 2개 이상의 multi-GPU 사용 (ex. 0번, 1번)

import os
os.environ['CUDA_VISIBLE_DEVICES'] = '0, 1'

 

 

이렇게 하고

 

1. Tensorflow

tensorflow.distribute의 MirroredStrategy를 바로 사용해주면 된다.

위에서 이미 GPU를 제한했기 때문에 따로 설정해줄 것은 없다.

단지 선언한 mirrored_strategy안에 모델 define부터 모델 생성, compile까지 다 들어가게 하면 된다.

 

 

import tensorflow as tf

 

mirrored_strategy = tf.distribute.MirroredStrategy()

with mirrored_strategy.scope():

     input = Input(input_shape=(30,))

     hidden_1 = Dense(20)(input)

     output = Dense(3)(hidden_1)

 

      model = Model(input, output)

      model.compile(loss='categorical_crossentropy', optimizer='adam')

 

 

2. pyTorch

pyTorch는 DataParallel을 이용하면 된다.

torch.device로 torch에 device를 할당해준 뒤 torch.nn.DataParallel을 이용하면 된다.

 

 

import torch

import torch.nn as nn

 

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

 

_net = '만들고자 하는 Network'

net = nn.DataParallel(_net).to(device)