Spaces:
Sleeping
Sleeping
File size: 2,980 Bytes
5194558 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 | #!/usr/bin/env python3
"""Train a tiny demo weather model and export it to ONNX.
This is intentionally simple: it learns from daily max temperature only and
predicts the next five max-temperature values from the previous seven days.
Use it as a scaffold for replacing the mock model with a stronger architecture.
"""
from __future__ import annotations
import sys
from pathlib import Path
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
import numpy as np
import pandas as pd
import torch
import torch.nn as nn
import torch.optim as optim
LOOKBACK_DAYS = 7
FORECAST_DAYS = 5
HISTORY_CSV_PATH = Path("data/tokyo_weather_history_30y.csv")
ONNX_PATH = Path("models/mock_model/model.onnx")
PT_PATH = Path("models/mock_model/model.pt")
class MovingAverageForecastModel(nn.Module):
def __init__(self, lookback: int, forecast: int) -> None:
super().__init__()
self.fc = nn.Linear(lookback, forecast)
def forward(self, x: torch.Tensor) -> torch.Tensor:
return self.fc(x)
def create_sequences(values: np.ndarray) -> tuple[torch.Tensor, torch.Tensor]:
x_values = []
y_values = []
for index in range(len(values) - LOOKBACK_DAYS - FORECAST_DAYS + 1):
x_values.append(values[index : index + LOOKBACK_DAYS])
y_values.append(values[index + LOOKBACK_DAYS : index + LOOKBACK_DAYS + FORECAST_DAYS])
return (
torch.tensor(np.array(x_values), dtype=torch.float32),
torch.tensor(np.array(y_values), dtype=torch.float32),
)
def load_training_values() -> np.ndarray:
if not HISTORY_CSV_PATH.exists():
raise FileNotFoundError(
f"{HISTORY_CSV_PATH} does not exist. Run scripts/fetch_historical_weather.py first."
)
history = pd.read_csv(HISTORY_CSV_PATH)
return history["high_c"].dropna().to_numpy(dtype=np.float32)
def main() -> None:
values = load_training_values()
x_train, y_train = create_sequences(values)
model = MovingAverageForecastModel(LOOKBACK_DAYS, FORECAST_DAYS)
optimizer = optim.Adam(model.parameters(), lr=0.01)
criterion = nn.MSELoss()
with torch.no_grad():
model.fc.weight.fill_(1.0 / LOOKBACK_DAYS)
model.fc.bias.fill_(0.0)
model.train()
for epoch in range(120):
optimizer.zero_grad()
predictions = model(x_train)
loss = criterion(predictions, y_train)
loss.backward()
optimizer.step()
ONNX_PATH.parent.mkdir(parents=True, exist_ok=True)
torch.save(model.state_dict(), PT_PATH)
dummy_input = torch.zeros(1, LOOKBACK_DAYS, dtype=torch.float32)
torch.onnx.export(
model,
dummy_input,
ONNX_PATH,
input_names=["input"],
output_names=["forecast"],
dynamic_axes={"input": {0: "batch"}, "forecast": {0: "batch"}},
opset_version=17,
)
print(f"Saved PyTorch weights to {PT_PATH}")
print(f"Exported ONNX model to {ONNX_PATH}")
if __name__ == "__main__":
main()
|