학습에 사용된 Slow Variable의 데이터 분포
해당 Slow Variable들의 Coupling Term의 데이터 분포
if __name__ == "__main__":
# 과거 시점 개수 선택 (2 또는 10)
use_memory_points = 10 # 여기서 2나 10으로 설정
fold = os.path.join(os.getcwd(), f'memory_models_v2_{use_memory_points}points') # 폴더명에 시점 수 추가
os.makedirs(fold, exist_ok=True)
model_save_path = os.path.join(fold, 'lstm_model.pth')
# 시스템 파라미터 설정
K = 8 # Number of globa-scale variables X
J = 32 # Number of local-scale Y variables per single global-scale X variable
F = 15.0 # Forcing
b = 10.0 # ratio of amplitudes
c = 10.0 # time-scale ratio
h = 1.0 # Coupling coefficient
# 평가 파라미터 설정
nt_pre = 20000 # Number of time steps for model spinup
nt = 20000 # Number of time steps
si = 0.005 # Sampling time interval
dt = 0.005 # Time step
dt = dt*2
si = si*2
noise = 0.03 # 원래 데이터에서는 noise를 추가했으나, noise가 성능에 영향을 주는 정도를 없애기 위해 이번에는 추가하지 않음
# Initial conditions
k = np.arange(K)
j = np.arange(J * K)
Xinit = s(k, K) * (s(k, K) - 1) * (s(k, K) + 1)
Yinit = 0 * s(j, J * K) * (s(j, J * K) - 1) * (s(j, J * K) + 1)
# Solving true model
X, Y, t2, _ = integrate_L96_2t_with_coupling(Xinit, Yinit, si, nt_pre+nt, F, h, b, c, dt=dt)
X = X[nt_pre:,:]
Y = Y[nt_pre:,:]
# Sub-sampling (tmeporal sparsity)
X_train = X[::2,:]
# First training routine where we target state at the next time-step
if use_memory_points == 2:
n_hist = 2 # 현재 시점과 2개의 과거 시점을 활용. 즉, X(t), X(t-2dt), X(t-4dt)
else:
n_hist = 10 # 현재 시점과 10개의 과거 시점을 활용. 즉, X(t), X(t-2dt), ..., X(t-20dt)
# 바로 다음 시점의 coupling term들을 예측. 즉, (X+2dt)
n_fut = 1
학습에 사용된 Slow Variable의 데이터 분포
해당 Slow Variable들의 Coupling Term의 데이터 분포
Xt = []
for i in range(2*n_hist+1+1):
Xt.append(X_train[i:-2*n_hist-2+i-n_fut+1,:])
Xt = np.transpose(np.array(Xt), (1, 0, 2)) # 학습 데이터
Xtpdt = X_train[2*n_hist+2+n_fut-1:,:] # 정답 데이터
Ndata = Xt.shape[0]
mu_X = np.zeros(X_train.shape[1])
sigma_X = np.max(np.abs(X_train), axis=0)
# 학습 데이터 준비 (float32로 변경)
Xt = torch.from_numpy(Xt).float()
Xtpdt = torch.from_numpy(Xtpdt).float()
sigma_X = torch.from_numpy(sigma_X).float()
mu_X = torch.from_numpy(mu_X).float()
# LSTM 모델 초기화를 float32로 변경
model = LSTMMemoryNetwork(
input_dim=8, # Lorenz 96의 K값
hidden_dim=256, # 은닉 차원
output_dim=8, # 출력 차원 (K와 동일)
num_layers=2, # LSTM 층 수
dropout=0.1 # 드롭아웃 비율
)
# 옵티마이저 설정
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
criterion = torch.nn.MSELoss()
# 학습 파라미터 수정
n_epochs = 50
batch_size = 512
# 학습 루프
model.train()
for epoch in tqdm(range(n_epochs)):
epoch_loss = 0
batch_count = 0
for i in range(0, len(Xt), batch_size):
batch_Xt = Xt[i:i+batch_size]
batch_y = Xtpdt[i:i+batch_size]
optimizer.zero_grad()
if use_memory_points == 2:
pred = stepper_2(batch_Xt, model, F, sigma_X, mu_X, dt)
else: # use_memory_points == 10
pred = stepper_10(batch_Xt, model, F, sigma_X, mu_X, dt)
loss = criterion(pred, batch_y)
loss.backward()
optimizer.step()
epoch_loss += loss.item()
batch_count += 1
avg_epoch_loss = epoch_loss / batch_count
print(f'Epoch {epoch}, Average Loss: {avg_epoch_loss:.10f}')
print(f'Training finished')