mirror of
https://github.com/webclinic017/drift.git
synced 2026-08-13 02:48:07 +00:00
19 lines
623 B
Python
19 lines
623 B
Python
|
from pytorch_forecasting import TemporalFusionTransformer
|
||
|
|
from pytorch_forecasting.metrics import QuantileLoss
|
||
|
|
|
||
|
|
def create_TemporalFusionTransformer(training_dataset, model_options):
|
||
|
|
# create the model
|
||
|
|
tft = TemporalFusionTransformer.from_dataset(
|
||
|
|
training_dataset,
|
||
|
|
learning_rate=0.03,
|
||
|
|
hidden_size=32,
|
||
|
|
attention_head_size=1,
|
||
|
|
dropout=0.1,
|
||
|
|
hidden_continuous_size=16,
|
||
|
|
output_size=7,
|
||
|
|
loss=QuantileLoss(),
|
||
|
|
log_interval=2,
|
||
|
|
reduce_on_plateau_patience=4
|
||
|
|
)
|
||
|
|
print(f"Number of parameters in network: {tft.size()/1e3:.1f}k")
|
||
|
|
return tft
|