How sennadir/ts_jepa works
TS-JEPA is a PyTorch time-series representation-learning project: it pretrains an encoder to predict masked patch embeddings, then reuses a saved encoder for downstream forecasting of the next patch. The principal path runs from CSV preparation and masking through encoder/predictor training and checkpointing, then through a frozen encoder and linear decoder to forecast metrics. The graph also retains the separate autoregressive evaluation entry point documented by the tree. Some evaluation wiring is inferred only from README descriptions and filenames; unsampled implementations are not expanded.
Main entry points: pretrain.py Pretrain loop · eval_forecast_last_pred.py Last-value forecast · encoder.py Time-series encoder · predictor.py JEPA predictor
Components 19
Data preparation
- Weather CSV
data/weather/weather.csv - Pretrain loader
src/data_loaders/data_loader.py - Masked dataset
src/data_loaders/data_class.py - Forecast dataset
src/data_loaders/data_class.py - Evaluation loader
src/data_loaders/data_loader.py
JEPA pretraining
- Pretrain loop
pretrain.py - EMA target encoder
src/models/encoder.py - JEPA predictor
src/models/predictor.py - Loss and EMA update
pretrain.py - Encoder checkpoint
logs/output_model/weather/lr_1e-07_ema_momentum_0.998_mask_ratio_0.7_ratio_patches_10_encoder_128_2_1_predictor_128_2_1_epoch_5000.pt
Forecast evaluation
- Last-value forecast
eval_forecast_last_pred.py - Autoregressive evaluation
eval_autoregressive.py - Linear decoder
src/models/decoder.py - Forecast metrics
main/utils.py
Model components
- Time-series encoder
src/models/encoder.py - Patch tokenizer
src/models/tokenizer.py - Transformer blocks
src/models/utils/modules.py - Mask operations
src/models/utils/mask_utils.py
Other components
- Researcher
Connections 23
- Researcher → Pretrain loop: starts pretraining
- Researcher → Last-value forecast: starts evaluation
- Researcher → Autoregressive evaluation: starts evaluation
- Pretrain loop → Pretrain loader: requests batches
- Pretrain loader → Masked dataset: wraps dataset
- Masked dataset → Weather CSV: reads CSV
- Pretrain loop → Time-series encoder: encodes patches
- Time-series encoder → Patch tokenizer: tokenizes patches
- Time-series encoder → Transformer blocks: applies blocks
- Pretrain loop → EMA target encoder: copies encoder
- Pretrain loop → JEPA predictor: predicts masked targets
- JEPA predictor → Transformer blocks: applies blocks
- JEPA predictor → Mask operations: selects tokens
- Pretrain loop → Mask operations: selects targets
- Pretrain loop → Loss and EMA update: computes loss and updates EMA
- Loss and EMA update → Encoder checkpoint: saves encoder
- Last-value forecast → Evaluation loader: requests windows
- Evaluation loader → Forecast dataset: wraps dataset
- Forecast dataset → Weather CSV: reads CSV
- Last-value forecast → Encoder checkpoint: loads weights
- Last-value forecast → Time-series encoder: encodes context
- Last-value forecast → Linear decoder: predicts next patch
- Last-value forecast → Forecast metrics: computes errors