Pytorch에서 GPU를 사용

device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
torch.tensor([1,2], device=device)
torch.tensor([1,2]).to(device)
torch.tensor([1,2]).cuda()

 

Argument에 따른 방식

import argparse
import torch

parser = argparse.ArgumentPasrser()
parser.add_argument('--cpu', action='store_true', help='run in cpu')
args = parser.parse_args()

if args.cpu:
    device = torch.device('cpu')
else:
    device = torch.device('cuda')

x = torch.tensor([1.,2.]).to(device)

 

특정 GPU

• shell 에서 환경 변수로 설정

>> export CUDA_VISIBLE_DEVICES = 1
>> python main.py

 

• in code

device = torch.device('cuda:1')

with torch.cuda.device(0):
    # GPU 0에 할당
    X = torch.tensor([1,2]).to("cuda")
    
    # GPU 1에 할당하도록 명시
    X = torch.tensor([1,2]).to(device)

 

 

 

 

 

 

 

 

https://y-rok.github.io/pytorch/2020/10/03/pytorch-gpu.html

Pytorch에서 GPU 사용하기