refactor: reorganize project structure by moving core modules and update import paths in API server
This commit is contained in:
@@ -0,0 +1,64 @@
|
||||
#!/usr/bin/env python
|
||||
# coding: utf-8
|
||||
import os
|
||||
import json
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.optim as optim
|
||||
from torch.utils.data import DataLoader
|
||||
from tqdm import tqdm
|
||||
from core.ndvi_data_loader import NDVITimeSeriesDataset
|
||||
|
||||
print("=" * 70)
|
||||
print("🚀 Training LSTM/GRU Time Series model for NDVI (REAL DATA & GPU)")
|
||||
print("=" * 70)
|
||||
|
||||
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
||||
print(f"[SYSTEM] Device: {device.type.upper()}")
|
||||
|
||||
dataset = NDVITimeSeriesDataset(sequence_length=3, spatial=False)
|
||||
dataloader = DataLoader(dataset, batch_size=4, shuffle=True)
|
||||
|
||||
class LSTMModel(nn.Module):
|
||||
def __init__(self, input_size=1, hidden_size=32, num_layers=2):
|
||||
super().__init__()
|
||||
self.lstm = nn.LSTM(input_size, hidden_size, num_layers, batch_first=True)
|
||||
self.fc = nn.Linear(hidden_size, 1)
|
||||
|
||||
def forward(self, x):
|
||||
out, _ = self.lstm(x)
|
||||
out = self.fc(out[:, -1, :]) # Take last time step
|
||||
return out
|
||||
|
||||
model = LSTMModel().to(device)
|
||||
criterion = nn.MSELoss()
|
||||
optimizer = optim.Adam(model.parameters(), lr=1e-3)
|
||||
|
||||
print("\n[TRAIN] Bắt đầu Training...")
|
||||
model.train()
|
||||
total_loss = 0
|
||||
for epoch in range(50):
|
||||
for inputs, targets in dataloader:
|
||||
inputs, targets = inputs.to(device), targets.to(device)
|
||||
optimizer.zero_grad()
|
||||
outputs = model(inputs)
|
||||
loss = criterion(outputs, targets)
|
||||
loss.backward()
|
||||
optimizer.step()
|
||||
total_loss += loss.item()
|
||||
|
||||
avg_loss = total_loss / (len(dataloader) * 50)
|
||||
print(f"✅ Training completed! Avg MSE Loss (RMSE): {avg_loss**0.5:.4f}")
|
||||
|
||||
model_dir = "ndvi_forecast_model"
|
||||
model_path = os.path.join(model_dir, "ndvi_lstm_real.pth")
|
||||
torch.save(model.state_dict(), model_path)
|
||||
print(f"\n[SAVE] Model saved to {model_path}")
|
||||
|
||||
with open(os.path.join(model_dir, "ndvi_lstm_real_info.json"), "w") as f:
|
||||
json.dump({
|
||||
"model_type": "LSTM Time Series (Real Data & GPU)",
|
||||
"target": "NDVI", "epoch": 50,
|
||||
"rmse": float(avg_loss**0.5), "mae": float(avg_loss)
|
||||
}, f, indent=2)
|
||||
print("[SAVE] Model info saved.")
|
||||
Reference in New Issue
Block a user