N-BEATS block which takes a basis function as an argument.
| 301 | 'Sigmoid'] |
| 302 | |
| 303 | class _NHITSBlock(nn.Module): |
| 304 | """ |
| 305 | N-BEATS block which takes a basis function as an argument. |
| 306 | """ |
| 307 | def __init__(self, n_time_in: int, n_time_out: int, n_x: int, |
| 308 | n_s: int, n_s_hidden: int, n_theta: int, n_mlp_units: list, |
| 309 | n_pool_kernel_size: int, pooling_mode: str, basis: nn.Module, |
| 310 | n_layers: int, batch_normalization: bool, dropout_prob: float, activation: str): |
| 311 | """ |
| 312 | """ |
| 313 | super().__init__() |
| 314 | |
| 315 | assert (pooling_mode in ['max','average']) |
| 316 | |
| 317 | n_time_in_pooled = int(np.ceil(n_time_in/n_pool_kernel_size)) |
| 318 | |
| 319 | if n_s == 0: |
| 320 | n_s_hidden = 0 |
| 321 | n_mlp_units = [n_time_in_pooled + (n_time_in+n_time_out)*n_x + n_s_hidden] + n_mlp_units |
| 322 | |
| 323 | self.n_time_in = n_time_in |
| 324 | self.n_time_out = n_time_out |
| 325 | self.n_s = n_s |
| 326 | self.n_s_hidden = n_s_hidden |
| 327 | self.n_x = n_x |
| 328 | self.n_pool_kernel_size = n_pool_kernel_size |
| 329 | self.batch_normalization = batch_normalization |
| 330 | self.dropout_prob = dropout_prob |
| 331 | |
| 332 | assert activation in ACTIVATIONS, f'{activation} is not in {ACTIVATIONS}' |
| 333 | activ = getattr(nn, activation)() |
| 334 | |
| 335 | if pooling_mode == 'max': |
| 336 | self.pooling_layer = nn.MaxPool1d(kernel_size=self.n_pool_kernel_size, |
| 337 | stride=self.n_pool_kernel_size, ceil_mode=True) |
| 338 | elif pooling_mode == 'average': |
| 339 | self.pooling_layer = nn.AvgPool1d(kernel_size=self.n_pool_kernel_size, |
| 340 | stride=self.n_pool_kernel_size, ceil_mode=True) |
| 341 | |
| 342 | hidden_layers = [] |
| 343 | for i in range(n_layers): |
| 344 | hidden_layers.append(nn.Linear(in_features=n_mlp_units[i], out_features=n_mlp_units[i+1])) |
| 345 | hidden_layers.append(activ) |
| 346 | |
| 347 | if self.batch_normalization: |
| 348 | hidden_layers.append(nn.BatchNorm1d(num_features=n_mlp_units[i+1])) |
| 349 | |
| 350 | if self.dropout_prob>0: |
| 351 | hidden_layers.append(nn.Dropout(p=self.dropout_prob)) |
| 352 | |
| 353 | output_layer = [nn.Linear(in_features=n_mlp_units[-1], out_features=n_theta)] |
| 354 | layers = hidden_layers + output_layer |
| 355 | |
| 356 | # n_s is computed with data, n_s_hidden is provided by user, if 0 no statics are used |
| 357 | if (self.n_s > 0) and (self.n_s_hidden > 0): |
| 358 | self.static_encoder = _StaticFeaturesEncoder(in_features=n_s, out_features=n_s_hidden) |
| 359 | self.layers = nn.Sequential(*layers) |
| 360 | self.basis = basis |