[Python, Pytorch] Tensor, Cuda 사용법

안녕하세요, 오늘은 Pytorch 내 Tensor와 Cuda 사용법에 대해서 알아보도록 하겠습니다.

 

Tensor 및 Cuda 사용법에 대한 내용은 다음과 같습니다.

 

• torch.tensor()
• torch.as_tensor()
• size()
• dtype
• device
• unsqueeze() & squeeze()
• permute() & transpose()

 

다음은 Tensor와 Cuda를 사용하기 전 Cuda와 관련된 상태를 알아보는 코드입니다.

 

import torch
print(torch.cuda.is_available())
print(torch.cuda.current_device())
print(torch.cuda.device_count())
print(torch.cuda.get_device_name(0))

위의 코드는 차례 대로 cuda 사용 가능 여부 (True or False), 현재 어떤 index의 cuda device로 설정(0), 사용할 수 있는 cuda device의 개수, 특정 index(0)의 cuda device의 name은 어떻게 되는 지 알수 있습니다.

 

---

 

다음은 Tensor를 초기화하고 dtype과 device를 설정 후 해당 Tensor의 size(), dtype, device를 출력해보는 코드에 대해서 알아보겠습니다.

 

import torch
device_cpu = torch.device('cpu')

tensor_x = torch.randint(0, 256, (3, 512, 512))
tensor_x = torch.as_tensor(tensor_x, dtype=torch.uint8, device=device_cpu)
# tensor_x = torch.as_tensor(tensor_x, dtype=torch.uint8).cpu()
# tensor_x = torch.as_tensor(tensor_x, dtype=torch.uint8).to(device=device_cpu)
print(tensor_x.size())
print(tensor_x.dtype)
print(tensor_x.device)

torch.Size([3, 512, 512])
torch.uint8
cpu

  

torch.randint() 함수를 통해 Tensor를 초기화 하고, torch.as_tensor() 함수를 통해 dtype과 device를 설정하였습니다.

해당 Tensor에 대해서 size(), dtype, device를 이용하여 출력한 결과는 위와 같습니다.

**device는 device=device_cpu, cpu(), to(device=device_cpu)의 형태로 설정할 수 있습니다.

 

device(cpu)가 아닌, device(cuda)로 설정하는 방법은 다음과 같습니다.

 

import torch
device_cuda = torch.device('cuda')

tensor_x = torch.randint(0, 256, (3, 512, 512))
tensor_x = torch.as_tensor(tensor_x, dtype=torch.uint8, device=device_cuda)
tensor_x = torch.as_tensor(tensor_x, dtype=torch.uint8).cuda()
tensor_x = torch.as_tensor(tensor_x, dtype=torch.uint8).to(device=device_cuda)
print(tensor_x.size())
print(tensor_x.dtype)
print(tensor_x.device)

---

이어서 Tensor를 unsqueeze(), squeeze(), permute(), transpose() 함수를 사용한 결과는 다음과 같습니다.

 

import torch

tensor_x = torch.randint(0, 256, (3, 512, 512))
tensor_x = torch.as_tensor(tensor_x, dtype=torch.uint8, device='cpu')
print("Before : ", tensor_x.size())

tensor_x = tensor_x.unsqueeze(0)
print("After : ", tensor_x.size())

Before :  torch.Size([3, 512, 512])
After :  torch.Size([1, 3, 512, 512])

 

최초에 (3, 512, 512)의 형태의 Tensor에 unsqueeze(0) 함수를 적용하면 0번째 인덱스에 size가 1인 차원이 생성됩니다.

이때, 0 --> k 로 수정하게 되면 k번째 인덱스에 size가 1인 차원이 생성됩니다.

 

import torch
device_cpu = torch.device('cpu')

tensor_x = torch.randint(0, 256, (3, 1, 512, 512))
tensor_x = torch.as_tensor(tensor_x, dtype=torch.uint8, device=device_cpu)
print("Before : ", tensor_x.size())

tensor_x = tensor_x.squeeze(1)
print("After : ", tensor_x.size())

Before :  torch.Size([3, 1, 512, 512])
After :  torch.Size([3, 512, 512])

최초에 (3, 1, 512, 512)의 형태에 Tensor 내 1번째 인덱스에 size가 1인 차원이 있는 상태에서 squeeze(1) 함수를 적용하게 되면 해당 인덱스의 size가 1인 차원이 제거됩니다.

 

import torch
device_cpu = torch.device('cpu')

tensor_x = torch.randint(0, 256, (3, 1, 512, 1, 512))
tensor_x = torch.as_tensor(tensor_x, dtype=torch.uint8, device=device_cpu)
print("Before : ", tensor_x.size())

tensor_x = tensor_x.squeeze()
print("After : ", tensor_x.size())

Before :  torch.Size([3, 1, 512, 1, 512])
After :  torch.Size([3, 512, 512])

 

위의 코드와 같이 size가 1인 차원이 2개 이상 존재하는 상황에서 squeeze()와 같이 인덱스 설정을 하지 않으면 size가 1인 모든 차원들이 제거됩니다.

 

import torch
device_cpu = torch.device('cpu')

tensor_x = torch.randint(0, 256, (3, 512, 512))
tensor_x = torch.as_tensor(tensor_x, dtype=torch.uint8, device=device_cpu)
print("Before : ", tensor_x.size())

tensor_x = tensor_x.permute(1, 2, 0)
print("After : ", tensor_x.size())

Before :  torch.Size([3, 512, 512])
After :  torch.Size([512, 512, 3])

permute() 함수는 차원의 순서를 바꿀 수 있으며, 위의 코드에서는 최초의 Tensor 차원에서 1번째 차원, 2번째 차원, 0번째 차원 순서대로 재 설정할 수 있습니다. 다음은 서로 다른 2개의 차원의 순서를 바꿀 수 있습니다.

 

import torch
device_cpu = torch.device('cpu')

tensor_x = torch.randint(0, 256, (3, 512, 512))
tensor_x = torch.as_tensor(tensor_x, dtype=torch.uint8, device=device_cpu)
print("Before : ", tensor_x.size())

tensor_x = tensor_x.transpose(0, 2).transpose(0, 1)
print("After : ", tensor_x.size())

Before :  torch.Size([3, 512, 512])
After :  torch.Size([512, 512, 3])

위의 코드는 transpose() 함수를 사용해서 최의 Tensor 차원에서 0번째, 2번째 차원의 순서를 조정하고 조정된 Tensor를 기준으로 0번째, 1번째 차원의 순서를 조정한 결과입니다.

 

지금까지, Pytorch 내 Tensor 및 Cuda를 사용법에 대해서 알아보았습니다.