[PyTorch] GPU에 텐서 할당하기

 

 

 

 

텐서를 생성할 때 device 옵션으로 cpu와 cuda(gpu)를 선택할 수 있다.

(아래 사진 참조, cpu와 cuda(gpu) 말고도 opengl, opencl 등등 더 많은 옵션을 선택할 수 있는 듯 하다.)

 

 

 

 

 

주석 보면서 이해하자.

import torch
cpu_tensor = torch.zeros(2, 3, device = "cpu") # cpu로 설정한 2행 3열 영행렬을 cpu_tensor라는 변수에 저장.
print(cpu_tensor)

device = torch.device("cuda")  # cuda(gpu)를 device라는 변수에 저장하고
gpu_tensor = cpu_tensor.to(device) # cpu에 올라간 텐서를 gpu로 옮기는게 to()다.
print(gpu_tensor)

 

 

 

 

핵심은 device = "cpu"를 활용해서 cpu에서 처리할 수 있기도 하고

to()나 cuda()를 통해서 gpu로 옮길 수있다.

cpu_tensor1 = cpu_tensor.to("cuda") # cpu위에 올라가있는 tensor를 gpu로 옮기는 첫 번째 방법.
print(cpu_tensor1)

cpu_tensor2 = cpu_tensor.cuda() # cpu위에 올라가있는 tensor를 gpu로 옮기는 두 번째 방법.
print(cpu_tensor2)

 

 

 

그리고 gpu에서 cpu로 옮기는건 역으로 생각하면 된다.

to()나 cpu()이다.

 

예시는 안들겠다.

검색하면 금방 나온다.