토니의 연습장

pytorch - torch.max(A, dim=0), torch.max(A, dim=1) 본문

기타

pytorch - torch.max(A, dim=0), torch.max(A, dim=1)

bellmake 2025. 8. 12. 14:18

 

1. dim=0 기준 : "행 기준/고정하고 열별로 값 찾기"

torch.max(A, dim=0) -> dim=0 : "dim=0(행) 기준으로 위아래에서 프레스로 눌러서 각 열별로 최대값만 살아남는다"
-> 각 열별로, 결과는 열의 개수만큼 값과 해당되는 행 인덱스 반환 -> 열 개수만큼 원소 갖는 1차원 형태 결과 나옴 (shape : 상수)
-> keepdim=True 를 해 주면 tensor(1행, 열개수만큼의 열) 의 2차원 행x열로 결과 나옴
    ex) torch.randn(3,4) 에서 torch.max(A, dim=0) 결과 원소 4개 갖는 1차원 행 나옴

          torch.max(A, dim=0, keepdim=True) 결과 tensor(1, 4) 형태 나옴

  • dim=0: 행(0축)을 따라 최댓값을 취해 행을 없애므로, 각 열별 최댓값을 반환.
    예) A.shape=(3,4) → values.shape=(4,), indices.shape=(4,), keepdim=True면 (1,4). 인덱스는 행 인덱스.

 

 

2. dim=1 기준 : "열 기준/고정하고 행별로 값 찾기"

torch.max(A, dim=1) -> dim=1 : "dim=1(열) 기준으로 양옆에서 프레스로 눌러서 각 행별로 최대값만 살아남는다"

-> 각 행별로, 결과는 행의 개수만큼 값과 해당되는 열 인덱스 반환  -> 행 개수만큼 원소 갖는 1차원 형태 결과 나옴 (shape : 상수)

-> keepdim=True 를 해 주면 tensor(행개수만큼의 행, 1열) 의 2차원 행x 열로 결과 나옴

    ex) torch.randn(3,4) 에서 torch.max(A, dim=1) 결과 원소 3개 갖는 1차원 행 나옴

          torch.max(A, dim=1, keepdim=True) 결과 tensor(3, 1) 형태 나옴

  • dim=1: 열(1축)을 따라 최댓값을 취해 열을 없애므로, 각 행별 최댓값을 반환.
    예) A.shape=(3,4) → values.shape=(3,), indices.shape=(3,), keepdim=True면 (3,1). 인덱스는 열 인덱스.