LSTM.py

import numpy as np

# --- 1. 샘플 데이터 및 하이퍼파라미터 정의 ---
input_dim = 3
hidden_dim = 4
output_dim = 2
sequence_length = 5
learning_rate = 0.01
epochs = 500

# 예: 시퀀스 데이터의 총 합이 특정 값보다 크면 1, 아니면 0으로 분류
sample_input = np.random.rand(sequence_length, input_dim)

if np.sum(sample_input) > (sequence_length * input_dim / 2):
    sample_y = np.array([1, 0]).reshape(-1, 1) # Class 0
else:
    sample_y = np.array([0, 1]).reshape(-1, 1) # Class 1
    
print(f"Sample Input Shape: {sample_input.shape}")
print(f"True Label: Class {np.argmax(sample_y)}")
print("-" * 30)


# --- 2. 필요 함수 정의 (활성화 함수 및 손실 함수) ---

# 시그모이드 활성화 함수
def sigmoid(x):
    return 1 / (1 + np.exp(-x))

# 시그모이드 함수의 도함수
def sigmoid_derivative(x):
    s = sigmoid(x)
    return s * (1 - s)

# 하이퍼볼릭 탄젠트(tanh) 활성화 함수
def tanh(x):
    return np.tanh(x)

# tanh 함수의 도함수
def tanh_derivative(x):
    return 1 - np.tanh(x)**2

# 소프트맥스 함수
def softmax(x):
    # 수치적 안정성을 위해 입력값에서 최댓값을 빼줌 (Overflow 방지)
    e_x = np.exp(x - np.max(x, axis=0, keepdims=True))
    return e_x / np.sum(e_x, axis=0, keepdims=True)

# 크로스 엔트로피 손실 함수
def cross_entropy_loss(y_pred, y_true):
    # y_pred에 아주 작은 값을 더해 log(0) 방지
    return -np.sum(y_true * np.log(y_pred + 1e-9))


# --- 3. NumpyLSTM 모델 클래스 ---

class NumpyLSTM:
    # 모델의 가중치와 파라미터를 초기화합니다.
    # - input_size: 입력 벡터의 차원
    # - hidden_size: 은닉 상태 및 셀 상태 벡터의 차원
    # - output_size: 출력 벡터(클래스 개수)의 차원
    def __init__(self, input_size, hidden_size, output_size, learning_rate=0.01):
        self.input_size = input_size
        self.hidden_size = hidden_size
        self.output_size = output_size
        self.learning_rate = learning_rate

        # LSTM 파라미터 초기화 (Forget, Input, Cell, Output 게이트)
        # 각 게이트는 입력(x)과 이전 은닉 상태(h)를 모두 받으므로, 가중치 행렬을 합쳐서 정의
        self.Wx = np.random.randn(4 * hidden_size, input_size) * 0.1
        self.Wh = np.random.randn(4 * hidden_size, hidden_size) * 0.1
        self.b = np.zeros((4 * hidden_size, 1))

        # Dense Layer (출력층) 파라미터 초기화
        self.Why = np.random.randn(output_size, hidden_size) * 0.1
        self.by = np.zeros((output_size, 1))
        
        # 그래디언트를 저장할 변수 초기화
        self.dWx, self.dWh, self.db = np.zeros_like(self.Wx), np.zeros_like(self.Wh), np.zeros_like(self.b)
        self.dWhy, self.dby = np.zeros_like(self.Why), np.zeros_like(self.by)

    # 순전파 과정을 수행합니다.
    # - inputs: (시퀀스 길이, 입력 차원) 형태의 2D numpy 배열
    # - y_true: (출력 차원, 1) 형태의 one-hot 인코딩된 정답 레이블
    def forward(self, inputs, y_true):
        self.inputs = inputs
        self.y_true = y_true
        seq_length = inputs.shape[0]

        # 이전 은닉 상태와 셀 상태를 저장할 딕셔너리
        self.h_states, self.c_states = {}, {}
        self.h_states[-1] = np.zeros((self.hidden_size, 1))
        self.c_states[-1] = np.zeros((self.hidden_size, 1))

        # 순전파에 필요한 중간 값들을 저장할 딕셔너리
        self.z_s, self.f_s, self.i_s, self.c_tilde_s, self.o_s = {}, {}, {}, {}, {}

        # 1. LSTM 셀 순전파 (시간 순서대로)
        for t in range(seq_length):
            xt = self.inputs[t].reshape(-1, 1) # 현재 타임스텝의 입력
            h_prev = self.h_states[t - 1]
            c_prev = self.c_states[t - 1]

            # (1) 게이트 계산을 위한 선형 결합
            # 4개의 게이트(f, i, c_tilde, o) 계산을 한 번에 수행
            self.z_s[t] = self.Wx @ xt + self.Wh @ h_prev + self.b
            
            # (2) 각 게이트 활성화
            # Forget Gate (망각 게이트)
            self.f_s[t] = sigmoid(self.z_s[t][:self.hidden_size, :])
            # Input Gate (입력 게이트)
            self.i_s[t] = sigmoid(self.z_s[t][self.hidden_size:2*self.hidden_size, :])
            # Cell Candidate (셀 상태 후보)
            self.c_tilde_s[t] = tanh(self.z_s[t][2*self.hidden_size:3*self.hidden_size, :])
            # Output Gate (출력 게이트)
            self.o_s[t] = sigmoid(self.z_s[t][3*self.hidden_size:, :])
            
            # (3) 셀 상태 및 은닉 상태 업데이트
            self.c_states[t] = self.f_s[t] * c_prev + self.i_s[t] * self.c_tilde_s[t]
            self.h_states[t] = self.o_s[t] * tanh(self.c_states[t])

        # 2. Dense Layer & Softmax 순전파
        self.final_h = self.h_states[seq_length - 1]
        self.logits = self.Why @ self.final_h + self.by
        self.y_pred = softmax(self.logits)
        
        # 3. 손실(Loss) 계산
        self.loss = cross_entropy_loss(self.y_pred, self.y_true)

        return self.loss, self.y_pred

    # 역전파(BPTT) 과정을 수행하여 그래디언트를 계산합니다.
    def backward(self):
        # 그래디언트 초기화
        self.dWx, self.dWh, self.db = np.zeros_like(self.Wx), np.zeros_like(self.Wh), np.zeros_like(self.b)
        self.dWhy, self.dby = np.zeros_like(self.Why), np.zeros_like(self.by)
        
        # 다음 타임스텝에서 넘어올 그래디언트 초기화
        dh_next = np.zeros_like(self.h_states[0])
        dc_next = np.zeros_like(self.c_states[0])

        # 1. Dense & Softmax Layer 역전파
        d_logits = self.y_pred - self.y_true # Loss에 대한 Logits의 그래디언트
        self.dWhy = d_logits @ self.final_h.T
        self.dby = d_logits
        dh_final = self.Why.T @ d_logits # LSTM의 최종 은닉 상태에 대한 그래디언트

        # dh_next에 최종 그래디언트 추가
        dh_next += dh_final

        # 2. LSTM 셀 역전파 (시간 역순으로)
        for t in reversed(range(len(self.inputs))):
            xt = self.inputs[t].reshape(-1, 1)
            h_prev = self.h_states[t - 1]
            c_prev = self.c_states[t - 1]
            
            # (1) 은닉 상태와 셀 상태에 대한 그래디언트 계산
            do = dh_next * tanh(self.c_states[t])
            dc = dc_next + dh_next * self.o_s[t] * tanh_derivative(self.c_states[t])
            
            # (2) 각 게이트의 활성화 이전 값(z)에 대한 그래디언트 계산
            dz_o = do * sigmoid_derivative(self.z_s[t][3*self.hidden_size:, :])
            dc_tilde = dc * self.i_s[t]
            dz_c = dc_tilde * tanh_derivative(self.z_s[t][2*self.hidden_size:3*self.hidden_size, :])
            di = dc * self.c_tilde_s[t]
            dz_i = di * sigmoid_derivative(self.z_s[t][self.hidden_size:2*self.hidden_size, :])
            df = dc * c_prev
            dz_f = df * sigmoid_derivative(self.z_s[t][:self.hidden_size, :])
            
            # (3) 4개의 그래디언트를 하나로 합치기
            dz = np.vstack((dz_f, dz_i, dz_c, dz_o))
            
            # (4) 파라미터에 대한 그래디언트 누적
            self.dWx += dz @ xt.T
            self.dWh += dz @ h_prev.T
            self.db += dz
            
            # (5) 이전 타임스텝으로 전달할 그래디언트 계산
            dh_next = self.Wh.T @ dz
            dc_next = self.f_s[t] * dc
            
        # 그래디언트 폭발(exploding gradients)을 방지하기 위한 클리핑
        for dparam in [self.dWx, self.dWh, self.db, self.dWhy, self.dby]:
            np.clip(dparam, -5, 5, out=dparam)

    # 계산된 그래디언트를 사용하여 파라미터를 업데이트합니다. (Gradient Descent)
    def update(self):
        self.Wx -= self.learning_rate * self.dWx
        self.Wh -= self.learning_rate * self.dWh
        self.b -= self.learning_rate * self.db
        self.Why -= self.learning_rate * self.dWhy
        self.by -= self.learning_rate * self.dby


# --- 4. 모델 학습 실행 ---
if __name__ == '__main__':
    # 모델 인스턴스 생성
    lstm = NumpyLSTM(input_size=input_dim, hidden_size=hidden_dim, output_size=output_dim, learning_rate=learning_rate)

    # 학습 루프
    for epoch in range(epochs):
        # 1. 순전파 (오타 수정됨)
        loss, y_pred = lstm.forward(sample_input, sample_y)
        
        # 2. 역전파
        lstm.backward()
        
        # 3. 가중치 업데이트
        lstm.update()

        if epoch % 100 == 0:
            print(f"Epoch {epoch}, Loss: {loss:.4f}")
            print(f"Predicted Probs: {y_pred.flatten()}")
            print(f"Predicted Class: {np.argmax(y_pred)}")
            print("-" * 20)

    print("\n--- Training Finished ---")
    final_loss, final_y_pred = lstm.forward(sample_input, sample_y)
    print(f"Final Loss: {final_loss:.4f}")
    print(f"Final Prediction: Class {np.argmax(final_y_pred)} (Probs: {final_y_pred.flatten()})")
    print(f"True Label: Class {np.argmax(sample_y)}")