프로그래밍/Python

텐서 차원 변경하기(pytorch - squeeze, unsqueeze/ numpy - squeeze, expand_dims)

히또아빠 2023. 12. 4. 09:57

텐서의 원소를 유지한채 모양과 차원크기 변환

  • 차원이 1인경우 해당 차원을 제거,
  • 특정 위치에 1인차원을 추가
  • dim은 제거하거나 추가할 차원의 위치

1.numpy - squeeze, expand_dims

# 1차원 배열 생성
array = np.array([[1], [2], [3]])

# 차원을 추가하여 2차원 배열로 만들기
expanded_array = np.expand_dims(array, axis=0)

# 차원을 제거하여 1차원 배열로 만들기
squeezed_array = np.squeeze(expanded_array, axis=0)

print("Original array:")
print(array)
print(array.shape)
print("\nAfter expand_dims:")
print(expanded_array)
print(expanded_array.shape)
print("\nAfter squeeze:")
print(squeezed_array)
print(squeezed_array.shape)


Original array:
[[1]
 [2]
 [3]]
(3, 1)

After expand_dims:
[[[1]
  [2]
  [3]]]
(1, 3, 1)

After squeeze:
[[1]
 [2]
 [3]]
(3, 1)

2.pytorch - squeeze, unsqeeze

# 1차원 텐서 생성
tensor = torch.tensor([[1], [2], [3]])

# 차원을 추가하여 2차원 텐서로 만들기
unsqueeze_tensor = torch.unsqueeze(tensor, dim=0)

# 차원을 제거하여 1차원 텐서로 만들기
squeeze_tensor = torch.squeeze(unsqueeze_tensor, dim=0)

print("Original tensor:")
print(tensor)
print(tensor.shape)
print("\nAfter unsqueeze:")
print(unsqueeze_tensor)
print(unsqueeze_tensor.shape)
print("\nAfter squeeze:")
print(squeeze_tensor)
print(squeeze_tensor.shape)

Original tensor:
tensor([[1],
        [2],
        [3]])
torch.Size([3, 1])

After unsqueeze:
tensor([[[1],
         [2],
         [3]]])
torch.Size([1, 3, 1])

After squeeze:
tensor([[1],
        [2],
        [3]])
torch.Size([3, 1])
300x250
반응형