[PYTORCH] nn.Linear의 입력 및 출력 차원 계산법 2가지와 텐서 셰이프 에러 해결 방법 7가지

nn.Linear

 

 

파이토치(PyTorch)를 활용해 딥러닝 아키텍처를 설계할 때 가장 기본이 되면서도 가장 잦은 에러를 유발하는 레이어가 바로 nn.Linear(선형 레이어)입니다. 많은 입문자가 단순히 "입력 뉴런 수와 출력 뉴런 수만 맞추면 된다"고 생각하지만, 실제 현업 모델에서는 다차원 배치 데이터, 시퀀스 데이터, 그리고 컨볼루션 레이어와의 연결 과정에서 복잡한 텐서 셰이프(Tensor Shape) 연산이 일어납니다. 이 메커니즘을 정확히 이해하지 못하면 RuntimeError: size mismatch라는 치명적인 장벽에 부딪히게 됩니다. 본 가이드에서는 시니어 딥러닝 엔지니어의 관점에서 nn.Linear의 선형 대수적 동작 원리와 다차원 텐서가 입력될 때의 차원 보존 법칙을 심층 분석하고, 실무에서 즉시 적용 가능한 7가지 텐서 셰이프 해결 전략을 제시합니다.

---

1. nn.Linear의 차원 계산법과 선형 변환의 차이

nn.Linear 레이어의 핵심은 $y = xA^T + b$ 연산입니다. 여기서 중요한 점은 파이토치가 입력 텐서의 마지막 차원만을 연산 대상으로 삼고, 그 앞의 차원들은 그대로 유지(Broadcasting)한다는 것입니다.

항목	설정 파라미터	입력 데이터 (Input)	출력 데이터 (Output)
표기법	nn.Linear(in_features, out_features)	(..., in_features)	(..., out_features)
데이터 형태	가중치 및 편향의 크기 결정	Batch, Sequence, Feature 차원 포함	Feature 차원만 변경됨
차원 유지	고정 (Static)	가변 (Dynamic)	입력의 앞 차원들을 그대로 보존
행렬 연산	$W: (out, in)$	$X: (Batch, in)$	$Y = XW^T + b: (Batch, out)$
에러 발생 시점	레이어 정의 시	Forward Pass 실행 시	in_features 불일치 시

---

2. 왜 다차원 텐서 계산법이 독창적인 가치를 지니는가?

• 자동 루프 처리: 파이토치의 nn.Linear는 입력이 $(N, L, H_{in})$인 3차원 텐서(예: NLP의 시퀀스 데이터)인 경우, 시퀀스 길이 $L$을 유지하면서 각 타임스텝의 은닉 상태 $H_{in}$에만 선형 변환을 적용합니다. 이는 별도의 반복문 없이 병렬 처리를 가능케 합니다.
• 유연한 전결합층 설계: 이미지 처리(CNN) 이후 Flatten을 거쳐 전결합층으로 넘어갈 때, 입력 차원의 총합을 정확히 계산하는 능력이 모델 전체 성능의 안정성을 결정합니다.
• 메모리 최적화: 선형 레이어의 입출력 차원을 정밀하게 설계함으로써 불필요한 파라미터 수를 줄이고 연산 속도를 극대화할 수 있습니다.

---

3. 실무 해결을 위한 텐서 셰이프 계산 예제 7가지

개발자가 실무 현장에서 마주하는 다양한 데이터 형태에 따라 nn.Linear의 차원을 맞추는 7가지 해결 시나리오입니다.

Example 1: 기본적인 2D 배치 데이터 처리 (MLP)

import torch
import torch.nn as nn

# 입력 특징 20개, 출력 특징 10개
fc = nn.Linear(in_features=20, out_features=10)
input_data = torch.randn(32, 20) # (Batch_size, in_features)
output = fc(input_data)
print(output.shape) # torch.Size([32, 10])
    

Example 2: NLP 시퀀스 데이터 처리 (3D 텐서)

마지막 차원만 변경되고 앞의 Batch와 Sequence 차원은 유지됩니다.

# (Batch, Seq_len, Hidden_dim) -> (Batch, Seq_len, Output_dim)
linear = nn.Linear(512, 128)
input_seq = torch.randn(16, 50, 512)
output_seq = linear(input_seq)
print(output_seq.shape) # torch.Size([16, 50, 128])
    

Example 3: CNN 출력과 Linear 연결 해결 (Flatten)

# Conv 출력: (Batch, Channels, Height, Width)
conv_out = torch.randn(8, 64, 7, 7)
# Linear에 넣기 위해 1차원으로 펼치기 (Batch 제외)
flattened = conv_out.view(conv_out.size(0), -1) 
# in_features = 64 * 7 * 7 = 3136
fc_layer = nn.Linear(3136, 1000)
output = fc_layer(flattened)
    

Example 4: 다중 차원(4D 이상) 텐서의 자동 매핑 해결

# (Batch, Depth, Height, Width, Features)
volumetric_data = torch.randn(4, 10, 32, 32, 16)
fc_deep = nn.Linear(16, 8)
# 마지막 차원 16이 8로 변환됨
output_vol = fc_deep(volumetric_data)
print(output_vol.shape) # torch.Size([4, 10, 32, 32, 8])
    

Example 5: LazyLinear를 활용한 입력 차원 자동 해결

PyTorch 1.10 이상에서 지원하는 방식으로, 첫 데이터 입력 시점에 in_features를 자동으로 추론합니다.

# 입력 차원을 명시하지 않고 정의
lazy_fc = nn.LazyLinear(out_features=256)
input_x = torch.randn(1, 4096)
# 첫 실행 시점에 차원이 4096으로 고정됨
output_y = lazy_fc(input_x)
    

Example 6: 가중치(Weight) 행렬에 직접 접근 및 차원 확인

linear_layer = nn.Linear(10, 5)
# 가중치 셰이프는 (out_features, in_features)로 저장됨에 주의
print(linear_layer.weight.shape) # torch.Size([5, 10])
print(linear_layer.bias.shape)   # torch.Size([5])
    

Example 7: 사이즈 불일치(Size Mismatch) 에러 실시간 디버깅

def check_linear_dim(layer, input_tensor):
    try:
        return layer(input_tensor)
    except RuntimeError as e:
        print(f"Dimension Mismatch! Input: {input_tensor.shape[-1]}, Expected: {layer.in_features}")
        # 동적으로 레이어 재정의 또는 입력 텐서 수정 로직 필요
    

---

4. 시니어의 인사이트: Linear 레이어와 메모리 레이아웃

실무에서 대규모 모델을 다룰 때 nn.Linear의 입력 차원이 매우 크면 행렬 곱셈 연산 시 메모리 대역폭 정체가 발생할 수 있습니다. 특히 텐서가 Contiguous(연속적)하지 않은 상태에서 view()나 reshape()을 거쳐 nn.Linear에 들어갈 경우 연산 효율이 급감합니다. 성능 최적화가 필요하다면 .contiguous() 호출 여부를 점검하고, 가능하다면 out_features를 하드웨어 가속기(GPU Tensor Core)에 최적화된 8의 배수로 설정하는 것이 좋습니다.

---

5. 결론 및 요약

파이토치 nn.Linear의 차원 계산은 단순한 수치 맞춤을 넘어 모델 구조의 유연성을 결정하는 핵심입니다.

• 핵심 원리: 입력 텐서의 마지막 차원만이 in_features와 일치해야 하며, 나머지 앞 차원들은 출력에서도 보존된다.
• 셰이프 해결: CNN과의 연결 시에는 view(-1) 또는 nn.Flatten을 통해 특징 차원을 일렬로 정렬해야 한다.
• 실무 팁: 다차원 배치를 다룰 때 nn.Linear가 제공하는 자동 브로드캐스팅 기능을 활용하면 코드의 복잡성을 획기적으로 낮출 수 있다.

---

참조 및 출처 (Sources)

• PyTorch Official Documentation: torch.nn.Linear .
• Deep Learning with PyTorch (Eli Stevens et al., Manning Publications): Chapter 8. Using Convolutions to Generalize
• Stanford CS231n: Linear Classification and Affine Layers