Numpy를 torch Tensor로 변환하기 [shape issue, numpy(), torch.from_numpy(), torch.Tensor()]

Segmentation을 위한 dataloader 구현 중 mask이미지를 numpy로 전처리한 후 transfroms.ToTensor()로 변환하니 shape가 바뀌는 것을 발견해서 정리하게 됐습니다

numpy 로 생성했던 shape: (3, 1024, 912)

transforms.ToTensor()적용 후 : (912, 3, 1024)

1. 이러한 문제가 발생한 이유 torchvision.transforms.ToTensor

ToTensor - Torchvision main documentation

Convert a or numpy.ndarray to tensor. This transform does not support torchscript. Converts a PIL Image or numpy.ndarray (H x W x C) in the range [0, 255] to a torch.FloatTensor of shape (C x H x W) in the range [0.0, 1.0] if the PIL Image belongs to one of the modes (L, LA, P, I, F, RGB, YCbCr, RGBA, CMYK, 1) or if the numpy.ndarray has dtype = np.uint8 In the other cases, tensors are returned without scaling.

https://pytorch.org/vision/main/generated/torchvision.transforms.ToTensor.html

pytorch doc를 살펴보면 PIL이미지나 numpy.ndarray를 tensor로 변환해주는 클래스로,

(HxWxC)형태의 [0,255]값을 갖는 PIL Image/numpy ndarray를 (C, H, W)의 shape와 [0.0, 1.0]의 값을 갖는 Tensor로 바꿔준다고 나와있습니다. 당연히 채널 순서가 바뀔수밖에 없네요.. torch tensor순서에 맞게 numpy array를 만들어줬기 때문에 이 클래스를 사용해서 Tensor로 변환했으면 안됐습니다 ㅎㅎ

저처럼 numpy array를 torch tensor 채널 순서에 맞춰 만든 경우 말고 일반적으로 PIL→ndarray→tensor순으로 변환할 때의 경우의 shape 변화를 정리해봤습니다.

PIL Image → ndarray → torch.tensor

2. torch.Tensor()

이 경우에도 ndarray와 shape는 변하지 않습니다 ㅎㅎ

그러나 이는 numpy array의 사본으로 취급되어 numpy array의 값은 변하지않고, 새로운 메모리를 할당하여 tensor를 생성합니다!

a=np.array([1,2,3,4])
b=torch.Tensor(a)
Python
복사

3. torch.from_numpy()

이 경우에도 ndarray와 shape는 변하지 않으나, ndarray의 dtype을 상속받으며 메모리버퍼를 공유해 tensor값이 변경될경우 numpy array의 값도 변경됩니다

a=np.array([1,2,3,4])
b=torch.from_numpy(a)
Python
복사

번외 : Tensor → numpy

import torch
import numpy as np

a = torch.rand(3,3)
b = a.numpy()
Python
복사