Linear Layer와 ReLU를 직접 구현해보기

2025-09-05 #pytorch #nn-module #relu

nn.Linear와 nn.ReLU 라는 이미 완성된 클래스가 존재하지만, 직접 구현해보면 더 명확하게 이해할 수 있고 재미있는 사실들도 발견할 수 있다. torch.nn 없이 직접 구현해보았다.

Linear Layer 구현하기#

선형 계층은 입력에 가중치를 곱하고 편향을 더하는 $y = W x + b$ 연산이다. torch.nn 모듈 없이 구현해보았다.

class WithoutNNLinear:
    def __init__(self, in_features, out_features):
        self.weight = torch.randn(size=(out_features, in_features))  # 가중치는 랜덤하게 초기화
        self.bias = torch.randn(size=(out_features,))               # 편향도 랜덤하게 초기화

    def __call__(self, x):
        return x @ self.weight.T + self.bias   # (..., in) @ (in, out) -> (..., out)

여기서 shape를 잘 봐야 한다.

weight의 shape는 (out_features, in_features)다.
bias의 shape는 (out_features,)다.

가중치는 왜

(in, out)이 아니라 (out, in)인가? (in, out)으로 저장하면 forward에서 x @ weight로 전치 없이 끝난다. 그런데도 PyTorch는 (out, in)을 쓰고 weight.T로 전치한다. 이유는 무엇을 한 행(dim=0)으로 묶느냐의 문제다.

(out, in) → 출력 뉴런 기준으로 묶기: weight[i]가 "i번 출력 뉴런이 받는 입력 가중치 벡터 전체", shape (in,)
(in, out) → 입력 기준으로 묶기: weight[i]가 "i번 입력이 모든 뉴런에 뿌리는 가중치"

신경망의 한 층은 곧 출력 뉴런들의 모음이고, 각 출력 뉴런이 독립적인 계산 단위다. 초기화(뉴런마다 fan-in 기준), 프루닝(뉴런 통째 제거), 해석("3번 뉴런이 뭘 학습했나") 전부 출력 뉴런 단위로 일어난다. 그래서 dim=0 인덱스가 곧 출력 뉴런 하나가 되도록 (out, in)으로 잡는다. row-major 텐서에서 각 뉴런의 가중치가 메모리에 연속으로 놓이는 이점도 있다.

수학 관습과도 맞는다. 선형변환은 보통 입력을 왼쪽에서 곱하는 $y = W x$ 로 쓰고, 이때 $W$ 가 자연스럽게 (out, in)이다.

전치 비용은 걱정 안 해도 된다. weight.T는 데이터를 복사하지 않는 view() 메서드이다. 실제 matmul 커널(F.linear)이 전치된 행렬을 그대로 받아 계산하도록 최적화돼 있다. nn.Linear도 내부적으로 x @ weight.T + bias를 쓴다.

ReLU 구현하기#

ReLU는 음수를 0으로 만들고 양수는 그대로 두는 함수다. 그만큼 구현이 정말 단순하며, 불린 인덱싱을 쓰면 한 줄로 끝난다.

class WithoutNNReLU:
    def __init__(self):
        pass

    def __call__(self, x):
        x[x < 0] = 0   # x의 모든 원소에 대해 0보다 작은 값은 0으로 만든다
        return x

x[x < 0] = 0이 핵심이다. x < 0은 x와 같은 모양의 불린 텐서(True/False)를 만들고, 그걸 인덱스로 쓰면 True인 위치, 즉 0보다 작거나 같은 원소의 인덱스에만 0을 대입한다. 반복문 없이 조건에 맞는 원소만 골라 한 번에 바꾸는 게 불린 인덱싱의 힘이다.

Linear Layer 구현하기#

ReLU 구현하기#

댓글