def _init_modules(self, input_sizes: Dict[str, int]):
self.lstms = nn.ModuleDict()
self.transfer_fcs = nn.ModuleDict()
self.heads = nn.ModuleDict()
self.dropout = nn.Dropout(p=self.cfg.output_dropout)
for idx, freq in enumerate(self._frequencies):
freq_input_size = input_sizes[freq]
if self._is_shared_mtslstm and idx > 0:
self.lstms[freq] = self.lstms[self._frequencies[
idx - 1]] # same LSTM for all frequencies.
self.heads[freq] = self.heads[self._frequencies[
idx - 1]] # same head for all frequencies.
else:
self.lstms[freq] = nn.LSTM(input_size=freq_input_size,
hidden_size=self._hidden_size[freq])
self.heads[freq] = get_head(self.cfg,
n_in=self._hidden_size[freq],
n_out=self.output_size)
if idx < len(self._frequencies) - 1:
for state in ["c", "h"]:
if self._transfer_mtslstm_states[state] == "linear":
self.transfer_fcs[f"{state}_{freq}"] = nn.Linear(
self._hidden_size[freq],
self._hidden_size[self._frequencies[idx + 1]])
elif self._transfer_mtslstm_states[state] == "identity":
self.transfer_fcs[f"{state}_{freq}"] = nn.Identity()
else:
pass
def __init__(self, cfg: Config):
super(CustomLSTM, self).__init__(cfg=cfg)
# in case this class is used for analysis of an EmbCudaLSTM, we need to initialize the embedding network
if cfg.model in ["embcudalstm", "lstm"] and cfg.embedding_hiddens:
self.embedding_net = FC(cfg=cfg)
self._has_embedding_net = True
input_size = len(cfg.dynamic_inputs) + cfg.embedding_hiddens[-1]
else:
self._has_embedding_net = False
input_size = len(cfg.dynamic_inputs + cfg.evolving_attributes +
cfg.static_attributes + cfg.hydroatlas_attributes)
if cfg.use_basin_id_encoding:
input_size += cfg.number_of_basins
self._hidden_size = cfg.hidden_size
self.cell = _LSTMCell(input_size=input_size,
hidden_size=self._hidden_size,
initial_forget_bias=cfg.initial_forget_bias)
self.dropout = nn.Dropout(p=cfg.output_dropout)
self.head = get_head(cfg=cfg,
n_in=self._hidden_size,
n_out=self.output_size)
def __init__(self, cfg: Config):
super(CudaLSTM, self).__init__(cfg=cfg)
if cfg.embedding_hiddens:
LOGGER.warning(
"## Warning: Embedding settings are ignored. Use EmbCudaLSTM for embeddings"
)
input_size = len(cfg.dynamic_inputs + cfg.static_inputs +
cfg.hydroatlas_attributes + cfg.camels_attributes)
if cfg.use_basin_id_encoding:
input_size += cfg.number_of_basins
if cfg.head.lower() == "umal":
input_size += 1
self.lstm = nn.LSTM(input_size=input_size, hidden_size=cfg.hidden_size)
self.dropout = nn.Dropout(p=cfg.output_dropout)
self.head = get_head(cfg=cfg,
n_in=cfg.hidden_size,
n_out=self.output_size)
self._reset_parameters()
def __init__(self, cfg: Dict):
super(Transformer, self).__init__(cfg=cfg)
# specify submodules of the model that can later be used for finetuning. Names must match class attributes
self.module_parts = ['embedding_net', 'encoder', 'head']
# embedding net before transformer
# this is necessary to ensure that the number of inputs into the self-attention layer
# is divisible by the number of heads
if not cfg.embedding_hiddens:
raise ValueError(
'Transformer requires config argument embedding_hiddens.')
if cfg.embedding_hiddens[-1] % cfg.transformer_nheads != 0:
raise ValueError(
"Embedding dimension must be divisible by number of transformer heads."
)
input_size = len(cfg.dynamic_inputs + cfg.static_attributes +
cfg.hydroatlas_attributes + cfg.evolving_attributes)
if cfg.use_basin_id_encoding:
input_size += cfg.number_of_basins
self.embedding_net = FC(cfg=cfg, input_size=input_size)
embedding_dim = cfg.embedding_hiddens[-1]
self._sqrt_embedding_dim = math.sqrt(embedding_dim)
# positional encoder
self.positional_encoding_type = cfg.transformer_positional_encoding_type
if self.positional_encoding_type.lower() == 'concatenate':
encoder_dim = embedding_dim * 2
elif self.positional_encoding_type.lower() == 'sum':
encoder_dim = embedding_dim
else:
raise RuntimeError(
f"Unrecognized positional encoding type: {self.positional_encoding_type}"
)
self.pos_encoder = PositionalEncoding(
embedding_dim=embedding_dim,
dropout=cfg.transformer_positional_dropout,
position_type=cfg.transformer_positional_encoding_type,
max_len=cfg.seq_length)
# positional mask
self._mask = None
# encoder
encoder_layers = nn.TransformerEncoderLayer(
d_model=encoder_dim,
nhead=cfg.transformer_nheads,
dim_feedforward=cfg.transformer_dim_feedforward,
dropout=cfg.transformer_dropout)
self.encoder = nn.TransformerEncoder(
encoder_layer=encoder_layers,
num_layers=cfg.transformer_nlayers,
norm=None)
# head (instead of a decoder)
self.dropout = nn.Dropout(p=cfg.output_dropout)
self.head = get_head(cfg=cfg, n_in=encoder_dim, n_out=self.output_size)
# init weights and biases
self._reset_parameters()
def __init__(self, cfg: Config):
super(EALSTM, self).__init__(cfg=cfg)
self._hidden_size = cfg.hidden_size
input_size_dyn = len(cfg.dynamic_inputs)
input_size_stat = len(cfg.static_inputs + cfg.camels_attributes +
cfg.hydroatlas_attributes)
if cfg.use_basin_id_encoding:
input_size_stat += cfg.number_of_basins
# If hidden units for a embedding network are specified, create FC, otherwise single linear layer
if cfg.embedding_hiddens:
self.input_net = FC(cfg=cfg)
else:
self.input_net = nn.Linear(input_size_stat, cfg.hidden_size)
# create tensors of learnable parameters
self.weight_ih = nn.Parameter(
torch.FloatTensor(input_size_dyn, 3 * cfg.hidden_size))
self.weight_hh = nn.Parameter(
torch.FloatTensor(cfg.hidden_size, 3 * cfg.hidden_size))
self.bias = nn.Parameter(torch.FloatTensor(3 * cfg.hidden_size))
self.dropout = nn.Dropout(p=cfg.output_dropout)
self.head = get_head(cfg=cfg,
n_in=cfg.hidden_size,
n_out=self.output_size)
# initialize parameters
self._reset_parameters()
def __init__(self, cfg: Dict):
super(Transformer, self).__init__(cfg=cfg)
# embedding net before transformer
self.embedding_net = InputLayer(cfg)
# ensure that the number of inputs into the self-attention layer is divisible by the number of heads
if self.embedding_net.output_size % cfg.transformer_nheads != 0:
raise ValueError(
"Embedding dimension must be divisible by number of transformer heads. "
"Use statics_embedding/dynamics_embedding and embedding_hiddens to specify the embedding."
)
self._sqrt_embedding_dim = math.sqrt(self.embedding_net.output_size)
# positional encoder
self._positional_encoding_type = cfg.transformer_positional_encoding_type
if self._positional_encoding_type.lower() == 'concatenate':
encoder_dim = self.embedding_net.output_size * 2
elif self._positional_encoding_type.lower() == 'sum':
encoder_dim = self.embedding_net.output_size
else:
raise RuntimeError(
f"Unrecognized positional encoding type: {self.positional_encoding_type}"
)
self.positional_encoder = _PositionalEncoding(
embedding_dim=self.embedding_net.output_size,
dropout=cfg.transformer_positional_dropout,
position_type=cfg.transformer_positional_encoding_type,
max_len=cfg.seq_length)
# positional mask
self._mask = None
# encoder
encoder_layers = nn.TransformerEncoderLayer(
d_model=encoder_dim,
nhead=cfg.transformer_nheads,
dim_feedforward=cfg.transformer_dim_feedforward,
dropout=cfg.transformer_dropout)
self.encoder = nn.TransformerEncoder(
encoder_layer=encoder_layers,
num_layers=cfg.transformer_nlayers,
norm=None)
# head (instead of a decoder)
self.dropout = nn.Dropout(p=cfg.output_dropout)
self.head = get_head(cfg=cfg, n_in=encoder_dim, n_out=self.output_size)
# init weights and biases
self._reset_parameters()
def __init__(self, cfg: Config):
super(GRU, self).__init__(cfg=cfg)
self.embedding_net = InputLayer(cfg)
self.gru = nn.GRU(input_size=self.embedding_net.output_size,
hidden_size=cfg.hidden_size)
self.dropout = nn.Dropout(p=cfg.output_dropout)
self.head = get_head(cfg=cfg,
n_in=cfg.hidden_size,
n_out=self.output_size)
def __init__(self, cfg: Config):
super(CudaLSTM, self).__init__(cfg=cfg)
self.embedding_net = InputLayer(cfg)
self.lstm = nn.LSTM(input_size=self.embedding_net.output_size,
hidden_size=cfg.hidden_size)
self.dropout = nn.Dropout(p=cfg.output_dropout)
self.head = get_head(cfg=cfg,
n_in=cfg.hidden_size,
n_out=self.output_size)
self._reset_parameters()
def __init__(self, cfg: Config):
super(CustomLSTM, self).__init__(cfg=cfg)
self.embedding_net = InputLayer(cfg)
self._hidden_size = cfg.hidden_size
self.cell = _LSTMCell(input_size=self.embedding_net.output_size,
hidden_size=self._hidden_size,
initial_forget_bias=cfg.initial_forget_bias)
self.dropout = nn.Dropout(p=cfg.output_dropout)
self.head = get_head(cfg=cfg,
n_in=self._hidden_size,
n_out=self.output_size)
def __init__(self, cfg: Config):
super(EALSTM, self).__init__(cfg=cfg)
self._hidden_size = cfg.hidden_size
self.embedding_net = InputLayer(cfg)
self.input_gate = nn.Linear(self.embedding_net.statics_output_size,
cfg.hidden_size)
# create tensors of learnable parameters
self.dynamic_gates = _DynamicGates(
cfg=cfg, input_size=self.embedding_net.dynamics_output_size)
self.dropout = nn.Dropout(p=cfg.output_dropout)
self.head = get_head(cfg=cfg,
n_in=cfg.hidden_size,
n_out=self.output_size)
def __init__(self, cfg: Config):
super(EmbCudaLSTM, self).__init__(cfg=cfg)
# embedding net before LSTM
if not cfg.embedding_hiddens:
raise ValueError('EmbCudaLSTM requires config argument embedding_hiddens.')
self.embedding_net = FC(cfg=cfg)
input_size = len(cfg.dynamic_inputs) + cfg.embedding_hiddens[-1]
self.lstm = nn.LSTM(input_size=input_size, hidden_size=cfg.hidden_size)
self.dropout = nn.Dropout(p=cfg.output_dropout)
self.head = get_head(cfg=cfg, n_in=cfg.hidden_size, n_out=self.output_size)
self._reset_parameters()
def __init__(self, cfg: Config):
super(GRU, self).__init__(cfg=cfg)
input_size = len(cfg.dynamic_inputs + cfg.static_inputs +
cfg.hydroatlas_attributes + cfg.camels_attributes)
if cfg.use_basin_id_encoding:
input_size += cfg.number_of_basins
if cfg.head.lower() == "umal":
input_size += 1
self.gru = nn.GRU(input_size=input_size, hidden_size=cfg.hidden_size)
self.dropout = nn.Dropout(p=cfg.output_dropout)
self.head = get_head(cfg=cfg,
n_in=cfg.hidden_size,
n_out=self.output_size)
def __init__(self, cfg: Config):
super(EALSTM, self).__init__(cfg=cfg)
self._hidden_size = cfg.hidden_size
input_size_stat = len(cfg.evolving_attributes + cfg.static_attributes + cfg.hydroatlas_attributes)
if cfg.use_basin_id_encoding:
input_size_stat += cfg.number_of_basins
# If hidden units for a embedding network are specified, create FC, otherwise single linear layer
if cfg.embedding_hiddens:
self.input_gate = FC(cfg=cfg)
else:
self.input_gate = nn.Linear(input_size_stat, cfg.hidden_size)
# create tensors of learnable parameters
self.dynamic_gates = _DynamicGates(cfg=cfg)
self.dropout = nn.Dropout(p=cfg.output_dropout)
self.head = get_head(cfg=cfg, n_in=cfg.hidden_size, n_out=self.output_size)
def __init__(self, cfg: Config):
super(LSTM, self).__init__(cfg=cfg)
input_size = len(cfg.dynamic_inputs + cfg.static_inputs +
cfg.camels_attributes + cfg.hydroatlas_attributes)
if cfg.use_basin_id_encoding:
input_size += cfg.number_of_basins
self._hidden_size = cfg.hidden_size
self.cell = _LSTMCell(input_size=input_size,
hidden_size=self._hidden_size,
initial_forget_bias=cfg.initial_forget_bias)
self.dropout = nn.Dropout(p=cfg.output_dropout)
self.head = get_head(cfg=cfg,
n_in=self._hidden_size,
n_out=self.output_size)
def __init__(self, cfg: Config):
super(ODELSTM, self).__init__(cfg=cfg)
if len(cfg.use_frequencies) < 2:
raise ValueError('ODELSTM needs at least two frequencies.')
if isinstance(cfg.dynamic_inputs, dict) or isinstance(
cfg.hidden_size, dict):
raise ValueError(
'ODELSTM does not support per-frequency input variables or hidden sizes.'
)
# Note: be aware that frequency_factors and slice_timesteps have a slightly different meaning here vs. in
# MTSLSTM. Here, the frequency_factor is relative to the _lowest_ (not the next-lower) frequency.
# slice_timesteps[freq] is the input step (counting backwards) in the next-*lower* frequency from where on input
# data at frequency freq is available.
self._frequency_factors = {}
self._slice_timesteps = {}
self._frequencies = sort_frequencies(cfg.use_frequencies)
self._init_frequency_factors_and_slice_timesteps()
# start to count the number of inputs
self.input_size = len(cfg.dynamic_inputs + cfg.static_attributes +
cfg.hydroatlas_attributes +
cfg.evolving_attributes)
if cfg.use_basin_id_encoding:
self.input_size += cfg.number_of_basins
if cfg.head.lower() == 'umal':
self.input_size += 1
self.lstm_cell = _LSTMCell(self.input_size, self.cfg.hidden_size,
cfg.initial_forget_bias)
self.ode_cell = _ODERNNCell(self.cfg.hidden_size,
self.cfg.hidden_size,
num_unfolds=self.cfg.ode_num_unfolds,
method=self.cfg.ode_method)
self.dropout = nn.Dropout(p=cfg.output_dropout)
self.head = get_head(cfg=cfg,
n_in=self.cfg.hidden_size,
n_out=self.output_size)
