[PyTorch] Matrix Multiplication

https://pytorch.org/blog/inside-the-matrix/

PyTorch의 matmul(), mm(), bmm()

PyTorch에서 행렬 곱셈(Matrix Multiplication)을 수행하는 대표적인 함수는 다음 세 가지임.

torch.matmul()
torch.mm()
torch.bmm()

 

세 함수 모두 행렬 곱셈을 수행하지만, 입력 Tensor의 차원과 Batch Dimension의 처리 방식에 차이가 있음.

개인적으로 범용으로 사용가능한
torch.matmul을 선호.

 

https://dsaint31.tistory.com/725

[LA] Matrix-Vector Multiplication

---

1. torch.matmul()

Signature

torch.matmul(input, other, *, out=None) → Tensor

 

Tensor Method 형태:

input.matmul(other)

 

out은 연산 결과를 저장할 출력 Tensor를 미리 지정하는 선택적 매개변수임.

• 기본값이 None이므로, 별도로 지정하지 않으면 PyTorch가 결과를 저장할 새로운 Tensor를 생성함.
• 지정할 경우, 해당 Tensor 객체의 메모리에 결과를 저장.
  • CPU Tensor 연산 결과를 CUDA Tensor에 저장하려는 경우는 적절하지 않음.
  • 가급적 결과 shape 를 만족하는 텐서를 out으로 해야 함.

다음은 간단한 실행 예제임:

A = torch.randn(2, 3)
B = torch.randn(3, 4)

C = torch.matmul(A, B)

 

반면 다음과 같이 출력 Tensor를 미리 지정할 수도 있음.

C = torch.empty(2, 4)

torch.matmul(A, B, out=C)

 

이 경우 연산 결과가 이미 존재하는 C에 기록됨.

즉, out은
결과를 저장할 메모리 공간을 지정하는 옵션임.

 

특징

 

matmul()은 입력 Tensor의 차원에 따라 동작이 달라지는 가장 널리 사용되는 행렬 곱셈 함수 임.

input 차원	other 차원	동작
1D	1D	벡터 내적
2D	2D	행렬 곱
1D	2D	벡터와 행렬의 곱
2D	1D	행렬과 벡터의 곱
ND	ND	Batch Matrix Multiplication + Broadcasting

---

1.1 1D × 1D: Vector Inner Product

내적으로 동작함.

a = torch.tensor([1., 2., 3.])
b = torch.tensor([4., 5., 6.])

y = torch.matmul(a, b)

print(y)

 

결과:

tensor(32.)

 

계산은 다음과 같음:

$$
1 \times 4 + 2 \times 5 + 3 \times 6 = 32
$$

• 1D Tensor끼리 matmul()을 수행하면 Vector Inner Product가 됨.

---

1.2 2D × 2D: Matrix Multiplication

A = torch.randn(2, 3)
B = torch.randn(3, 4)

C = torch.matmul(A, B)

print(C.shape)

 

결과:

torch.Size([2, 4])

수학적으로는 다음과 같음:

$$
(2 \times 3)(3 \times 4) = (2 \times 4)
$$

---

1.3 3D × 3D: Batch Matrix Multiplication

뒤의 2개 축을 제외한 나머지들은 batch를 구성!

A = torch.randn(10, 2, 3)
B = torch.randn(10, 3, 4)

C = torch.matmul(A, B)

print(C.shape)

 

결과:

torch.Size([10, 2, 4])

 

각 Batch에 대해 독립적으로 행렬 곱셈을 수행함:

$$
C_i = A_iB_i
$$

즉, 다음이 성립:

A : (10, 2, 3)
B : (10, 3, 4)
C : (10, 2, 4)

---

1.4 matmul()의 Broadcasting

matmul()은

행렬 차원(Matrix Dimensions)을 제외한 Batch Dimensions에 일반적인 PyTorch Broadcasting 규칙을 적용함.

 

다음의 예를 참고:

A = torch.randn(10, 2, 3)
B = torch.randn(3, 4)

C = torch.matmul(A, B)

print(C.shape)

 

결과:

torch.Size([10, 2, 4])

 

B에는 Batch Dimension이 없으므로 모든 Batch에 공통으로 사용됨.

 

개념적으로는 다음과 같음.

for i in range(10):
    C[i] = A[i] @ B

• 실제로 Python Loop를 수행한다는 의미는 아님(같은 결과를 내지만, 병렬로 처리됨)
• Broadcasting을 통해 동일한 행렬이 각 Batch에 적용됨.

Batch의 Dimension 이 Broadcasting이 안되는 shape인 경우엔 동작안함.

 

다음의 예를 실행해볼 것:

A = torch.randn(10, 2, 3)
B = torch.randn(5, 3, 4)

C = torch.matmul(A, B)

 

batch의 10과 5는 Broadcasting 규칙을 만족하지 않음.

matmul()은
Batch Dimension에 Broadcasting을 지원하지만,
Broadcasting 규칙을 만족하는 경우에만 적용됨.

---

2. torch.mm()

Signature

torch.mm(input, mat2, *, out=None) → Tensor

특징

torch.mm()은 2D Tensor와 2D Tensor의 행렬 곱만 지원함.

 

입력 형태:

(m, n) @ (n, p)

 

출력 형태:

(m, p)

 

즉, 다음의 수식을 구현하고 있음:

$$
(m \times n)(n \times p) = (m \times p)
$$

 

기본 사용법

A = torch.randn(2, 3)
B = torch.randn(3, 4)

C = torch.mm(A, B)

print(C.shape)

 

결과:

torch.Size([2, 4])

---

mm()은 Broadcasting을 지원하지 않음

mm()은 2D × 2D만 허용함

• Batch Dimension을 포함하는 Tensor는 사용할 수 없음.

다음과 같은 ndim=3인 경우는 아예 동작안함:

A = torch.randn(10, 2, 3)
B = torch.randn(3, 4)

torch.mm(A, B)

mm()은
Batch Matrix Multiplication을 수행하지 않음.

---

3. torch.bmm()

Signature

torch.bmm(input, mat2, *, out=None) → Tensor

 

특징

torch.bmm()은 3D Tensor끼리의 Batch Matrix Multiplication을 수행함.

 

입력 형태:

(b, m, n) @ (b, n, p)

 

출력 형태:

(b, m, p)

 

즉, 다음의 수식을 구현하고 있음:

$$
(b,m,n)(b,n,p) \rightarrow (b,m,p)
$$

기본 사용법

A = torch.randn(10, 2, 3)
B = torch.randn(10, 3, 4)

C = torch.bmm(A, B)

print(C.shape)

 

결과:

torch.Size([10, 2, 4])

 

Batch의 각 요소별로 독립적인 행렬 곱을 수행.

C[0] = A[0] @ B[0]
C[1] = A[1] @ B[1]
...
C[9] = A[9] @ B[9]

---

bmm()은 Broadcasting을 지원하지 않음

다음의 예를 실행해 볼 것:

A = torch.randn(10, 2, 3)
B = torch.randn(1, 3, 4)

 

Batch Dimension은 다음과 같음.

A : (10, 2, 3)
B : ( 1, 3, 4)

10과 1은 일반적인 Broadcasting 규칙에 따라 Broadcasting 가능함.

 

하지만, bmm은 아예 Broadcasting을 지원하지 않으므로 동작안함:

C = torch.bmm(A, B)

• 오류가 발생함.

bmm()은
3D Batch Matrix Multiplication을 수행하지만,
Batch Dimension에 Broadcasting을 적용하지 않음.

---

4. 세 함수의 비교

함수	입력 형태	Broadcasting	주요 용도
torch.matmul()	1D, 2D, ND	Batch Dimension에 지원	범용 행렬 곱
torch.mm()	2D × 2D	지원하지 않음	일반 행렬 곱
torch.bmm()	3D × 3D	지원하지 않음	Batch 행렬 곱

---

5. 여러 Example

일반 행렬 곱

A = torch.randn(2, 3)
B = torch.randn(3, 4)

torch.matmul(A, B)  # 가능
torch.mm(A, B)      # 가능
torch.bmm(A, B)     # 불가능

 

Batch 행렬 곱

A = torch.randn(10, 2, 3)
B = torch.randn(10, 3, 4)

torch.matmul(A, B)  # 가능
torch.mm(A, B)      # 불가능
torch.bmm(A, B)     # 가능

 

Broadcasting 가능한 Batch Dimension

A = torch.randn(10, 2, 3)
B = torch.randn(1, 3, 4)

torch.matmul(A, B)  # 가능
torch.bmm(A, B)     # 불가능

---

참고: @ 연산자와의 관계

Python의 @ 연산자는 Tensor의 행렬 곱 연산을 수행함.

 

정확히 말하면, @는
객체의 __matmul__() 메서드에 연결되는 연산자 오버로딩(operator overloading) 된 연산자임
(Python 3.5 이상부터)

• Python의 기본 자료형인 int, float, list 등은 __matmul__() 메서드를 구현하고 있지않으나
• Torch의 tensor나 Numpy의 array 에는 구현되어 있음.

PyTorch의 tensor객체 A와 B에 대해 다음을 수행할 경우:

C = A @ B

 

이는 일반적으로 다음과 같은 의미임:

C = torch.matmul(A, B)

 

예를 들면 다음과 같음:

A = torch.randn(10, 2, 3)
B = torch.randn(1, 3, 4)

C1 = A @ B
C2 = torch.matmul(A, B)

print(torch.equal(C1, C2))

 

결과:

True

---

같이보면 좋은 자료들

2025.03.07 - [Python] - [DL] Tensor 다루기 기초 - PyTorch 중심

[DL] Tensor 다루기 기초 - PyTorch 중심

https://gist.github.com/dsaint31x/f99b09f2174fecf93d4643aff5b40204