📌 JAX로 시계열 예측 모델 구현 - RNN/LSTM 기반 미래 데이터 예측

📌 JAX로 시계열 예측 모델 구현 - RNN/LSTM 기반 미래 데이터 예측

---

🚀 시계열 예측이란?

시계열 예측은 시간 순서로 정렬된 데이터를 기반으로 미래 값을 예측하는 기법입니다.
기온 예측, 주식 가격 예측, 수요 예측 등 다양한 분야에서 활용됩니다.

이번 글에서는 JAX를 활용하여 **RNN(Recurrent Neural Network)**과 LSTM(Long Short-Term Memory) 기반의 시계열 예측 모델을 구현합니다.

---

💡 1. 시계열 데이터의 특징

• 순차성(Sequential): 이전 값이 다음 값에 영향을 줌
• 패턴 반복: 주기성, 계절성
• 불안정성: 노이즈와 이상치 존재 가능성

---

🔧 2. 라이브러리 설치

pip install jax jaxlib flax optax pandas matplotlib scikit-learn

---

💾 3. 데이터 준비 - 일별 기온 예측 예제

📥 CSV 또는 시계열 데이터셋 사용

import pandas as pd
from sklearn.preprocessing import MinMaxScaler
import jax.numpy as jnp

# 예시: 서울 기온 데이터 (일자, 평균기온)
df = pd.read_csv('seoul_temperature.csv')
temps = df['avg_temp'].values.reshape(-1, 1)

# 정규화
scaler = MinMaxScaler()
temps_scaled = scaler.fit_transform(temps)

# 시계열 샘플 생성
def create_sequences(data, window=30):
    xs, ys = [], []
    for i in range(len(data) - window):
        xs.append(data[i:i+window])
        ys.append(data[i+window])
    return jnp.array(xs), jnp.array(ys)

x_data, y_data = create_sequences(temps_scaled)

---

🧠 4. RNN 모델 구현 (Flax 기반)

반응형

📝 RNNCell 정의

from flax import linen as nn

class SimpleRNNCell(nn.Module):
    hidden_size: int

    @nn.compact
    def __call__(self, carry, x):
        h = carry
        h = nn.tanh(nn.Dense(self.hidden_size)(x) + nn.Dense(self.hidden_size)(h))
        return h, h

---

🔁 RNN 전체 모델 정의

class RNNModel(nn.Module):
    hidden_size: int
    output_size: int = 1

    @nn.compact
    def __call__(self, x):
        batch_size, seq_len, _ = x.shape
        rnn_cell = SimpleRNNCell(self.hidden_size)
        h = jnp.zeros((batch_size, self.hidden_size))

        for t in range(seq_len):
            h, _ = rnn_cell(h, x[:, t, :])

        output = nn.Dense(self.output_size)(h)
        return output

---

📉 5. 손실 함수 및 학습 루프

⚙️ 훈련 설정

import optax
from flax.training import train_state

model = RNNModel(hidden_size=64)
key = jax.random.PRNGKey(42)
params = model.init(key, jnp.ones((1, 30, 1)))  # input: batch, seq, feature

tx = optax.adam(1e-3)
state = train_state.TrainState.create(apply_fn=model.apply, params=params, tx=tx)

---

🧮 손실 및 학습

@jax.jit
def mse_loss(params, x, y):
    preds = model.apply(params, x)
    return jnp.mean((preds.squeeze() - y.squeeze()) ** 2)

@jax.jit
def train_step(state, x, y):
    loss, grads = jax.value_and_grad(mse_loss)(state.params, x, y)
    state = state.apply_gradients(grads=grads)
    return state, loss

---

🔁 6. 모델 학습

batch_size = 64
epochs = 20

for epoch in range(epochs):
    for i in range(0, len(x_data), batch_size):
        x_batch = x_data[i:i+batch_size]
        y_batch = y_data[i:i+batch_size]
        state, loss = train_step(state, x_batch[..., None], y_batch)
    print(f"Epoch {epoch+1}, Loss: {loss:.4f}")

---

✅ 7. 미래 값 예측 및 시각화

import matplotlib.pyplot as plt

def predict_future(state, input_seq, steps=10):
    results = []
    current = input_seq
    for _ in range(steps):
        pred = model.apply(state.params, current[None, ..., None])
        results.append(pred.squeeze())
        current = jnp.concatenate([current[1:], pred.squeeze()[None]], axis=0)
    return jnp.array(results)

# 예측
last_seq = x_data[-1]
future = predict_future(state, last_seq, steps=30)
future_rescaled = scaler.inverse_transform(future.reshape(-1, 1))

# 시각화
plt.plot(range(len(temps)), temps, label='원본')
plt.plot(range(len(temps), len(temps) + 30), future_rescaled, label='예측', color='red')
plt.legend()
plt.title("RNN 기반 기온 예측")
plt.show()

---

📌 다음 글 예고: JAX로 Attention 기반 시계열 예측 모델 구현 (Transformer for Time Series)

다음 글에서는 Self-Attention 구조를 적용한 시계열 예측 모델을 JAX로 구현하여,
복잡한 주기성과 장기 의존성까지 반영한 예측을 실습해보겠습니다.

---

 

JAX, 시계열 예측, RNN, LSTM, 기온 예측, 딥러닝, Time Series, Python, 고속 연산, 학습 루프, Flax, 실전 모델, 미래 예측, 데이터 전처리, 인공지능, 고성능 계산, 시각화