본문 바로가기
함수 설명/인공지능 (Pytorch)

[PyTorch] 텐서 펼치기: torch.flatten() 사용 가이드

by First Adventure 2024. 7. 19.
반응형

소개

  torch.flatten 함수는 PyTorch에서 입력 텐서를 1차원 텐서로 변환하는 함수입니다. 이 함수는 신경망의 마지막 단계에서 다차원 데이터를 평탄화(flatten)하여 완전 연결층(fully connected layer)으로 입력할 때 주로 사용됩니다.

 

기본 사용법

상세 설명

  • torch.flatten 함수의 기본 구문은 torch.flatten(input, start_dim=0, end_dim=-1)입니다.
    • input: 평탄화할 입력 텐서.
    • start_dim: 평탄화를 시작할 차원. 기본값은 0입니다.
    • end_dim: 평탄화를 종료할 차원. 기본값은 -1이며, 마지막 차원을 의미합니다.
  • torch.flatten 함수는 입력 텐서의 지정된 차원부터 마지막 차원까지를 하나의 1차원 텐서로 평탄화합니다.

예시 설명

  • torch.flatten(tensor)는 2차원 텐서 [[1, 2], [3, 4]]를 1차원 텐서 [1, 2, 3, 4]로 평탄화합니다.
  • torch.flatten(tensor, start_dim=1)는 3차원 텐서 [[[1, 2], [3, 4]], [[5, 6], [7, 8]]]를 2차원 텐서 [[1, 2, 3, 4], [5, 6, 7, 8]]로 평탄화합니다.
  • torch.flatten(tensor_gpu)는 GPU 장치에서 2차원 텐서 [[1, 2], [3, 4]]를 1차원 텐서 [1, 2, 3, 4]로 평탄화합니다.
# 기본
import torch

tensor = torch.tensor([[1, 2], [3, 4]])
flattened = torch.flatten(tensor)
print(flattened)
# 출력: tensor([1, 2, 3, 4])

# 다양한 크기와 데이터 타입
import torch

# 시작 차원과 종료 차원을 지정하여 평탄화
tensor = torch.tensor([[[1, 2], [3, 4]], [[5, 6], [7, 8]]])
flattened = torch.flatten(tensor, start_dim=1)
print(flattened)
# 출력: tensor([[1, 2, 3, 4],
#              [5, 6, 7, 8]])

# GPU에서 텐서 평탄화
tensor_gpu = torch.tensor([[1, 2], [3, 4]], device=torch.device('cuda'))
flattened_gpu = torch.flatten(tensor_gpu)
print(flattened_gpu)
# 출력: tensor([1, 2, 3, 4], device='cuda:0')

 

라이센스

  PyTorch의 표준 라이브러리와 내장 함수들은 BSD-style license 하에 배포됩니다. 이 라이센스는 자유 소프트웨어 라이센스로, 상업적 사용을 포함한 거의 모든 용도로 사용이 가능합니다. 라이센스와 저작권 정보는 PyTorch의 공식 GitHub 리포지토리에서 확인할 수 있습니다.

 

관련 내용

  [PyTorch] 차원을 추가하라: torch.unsqueeze() 함수 사용법 이해하기

  [PyTorch] 차원을 줄여라: torch.squeeze() 함수 사용법 이해하기

  [PyTorch] 차원을 추가하여 텐서 쌓기: torch.stack() 함수 활용법

  [PyTorch] 텐서 이어 붙이기: torch.cat() 함수 활용법

반응형