wassname/DeepTime
1
0
Fork 0
mirror of https://github.com/wassname/DeepTime.git synced 2026-08-07 11:16:49 +08:00
Code Issues Packages Projects Releases Wiki Activity

partially converted to M2S (multi inputs)

Browse Source
This commit is contained in:
wassname
2022-11-22 16:42:58 +08:00
parent 00ae1d8b8f
commit 9072613d03
7 changed files with 647 additions and 38 deletions
Download Patch File Download Diff File Expand all files Collapse all files
models/__init__.py
+5 -2
View File
@@ -4,13 +4,16 @@ import torch
from .DeepTIMe import deeptime
from .DeepTIMe2 import deeptime2
from .DeepTIMe3 import deeptime3
def get_model(model_type: str, **kwargs: Union[int, float]) -> torch.nn.Module:
if model_type == 'deeptime':
model = deeptime(datetime_feats=kwargs['datetime_feats'])
model = deeptime(datetime_feats=kwargs['datetime_feats'], dim_size=kwargs['dim_size'])
elif model_type=="deeptime2":
model = deeptime2(datetime_feats=kwargs['datetime_feats'])
model = deeptime2(datetime_feats=kwargs['datetime_feats'], dim_size=kwargs['dim_size'])
elif model_type=="deeptime3":
model = deeptime3(datetime_feats=kwargs['datetime_feats'], dim_size=kwargs['dim_size'])
else:
raise ValueError(f"Unknown model type {model_type}")
return model