os.environ["CUDA_VISIBLE_DEVICES"]= 사용하고자 하는 GPU
이때 사용하고자 하는 GPU는
0번 GPU의 경우 "0"
1번 GPU의 경우 "1"
0번 1번을 같이 쓰고 싶을 때는 "0,1"
이런 식으로 지정해주면 된다.
이렇게도 안되는 경우 ... (사실 이렇게 하면 그냥 다 되는 것 같다)
GPU_NUM = 0 # 원하는 GPU 번호 입력
device = torch.device(f'cuda:{GPU_NUM}' if torch.cuda.is_available() else 'cpu')
torch.cuda.set_device(device) # change allocation of current GPU여기까지가 이제 여러 개의 GPU가 있을 때 하나의 GPU만 사용하고 싶을 때 사용하는 방법이다.
여기서 부터는 여러개의 GPU 중 특정 여러개의 GPU를 사용하고 싶을 때 사용하는 방법이다. (N개의 GPU가 있을 때 N-1개 이하의 GPU를 사용하고 싶을 경우)
torch.nn.DataParallel
model = torch.nn.DataParallel(model, device_ids=[0,2])이렇게 하면 만약 0,1,2번의 GPU가 있다고 하면 0,2번의 GPU만을 사용하게 된다.
그리고 이 뒤로 model에 접근할 때는 model.module.#
여기서 #이 model에서 접근할 부분
만약 device_ids를 설정해주지 않으면 다 잡어버린다.
예를들어 freeze_bn()을 사용할 때
model.freeze_bn() # 이 아닌
model.module.freeze_bn() # .module 을 설정해주어야 한다.즉 DataParallel을 킬 경우 해당 .module을 붙여줘야한다.
