google-research

1
# coding=utf-8
2
# Copyright 2024 The Google Research Authors.
3
#
4
# Licensed under the Apache License, Version 2.0 (the "License");
5
# you may not use this file except in compliance with the License.
6
# You may obtain a copy of the License at
7
#
8
#     http://www.apache.org/licenses/LICENSE-2.0
9
#
10
# Unless required by applicable law or agreed to in writing, software
11
# distributed under the License is distributed on an "AS IS" BASIS,
12
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13
# See the License for the specific language governing permissions and
14
# limitations under the License.
15

16
"""Package for dealing with different models."""
17

18
from typing import Any, Dict
19

20
from models import base
21
from models import lstm_seq2seq
22
from models import lstm_seq2seq_saf
23
from models import tft
24
from models import tft_saf
25
import tensorflow as tf
26

27

28
def get_model_type(model_type, chosen_hparams,
29
                   loss_form):
30
  """Return a forecast model based on the type."""
31
  if loss_form == "MAE":
32
    training_loss_object = tf.keras.losses.MeanAbsoluteError(
33
        name="training_loss")
34
    self_supervised_loss_object = tf.keras.losses.MeanAbsoluteError(
35
        name="self_supervised_loss")
36
  elif loss_form == "MSE":
37
    training_loss_object = tf.keras.losses.MeanSquaredError(
38
        name="training_loss")
39
    self_supervised_loss_object = tf.keras.losses.MeanSquaredError(
40
        name="self_supervised_loss")
41
  else:
42
    raise Exception("The loss type is not supported")
43

44
  if model_type == "lstm_seq2seq":
45
    model = lstm_seq2seq.ForecastModel(
46
        loss_object=training_loss_object, hparams=chosen_hparams)
47
  elif model_type == "lstm_seq2seq_saf":
48
    model = lstm_seq2seq_saf.ForecastModel(
49
        loss_object=training_loss_object,
50
        self_supervised_loss_object=self_supervised_loss_object,
51
        hparams=chosen_hparams)
52
  elif model_type == "tft":
53
    model = tft.ForecastModel(
54
        loss_object=training_loss_object, hparams=chosen_hparams)
55
  elif model_type == "tft_saf":
56
    model = tft_saf.ForecastModel(
57
        loss_object=training_loss_object,
58
        self_supervised_loss_object=self_supervised_loss_object,
59
        hparams=chosen_hparams)
60
  else:
61
    raise Exception("The chosen model is not supported")
62

63
  return model
64