Residual 1D block with temporal convolutions. Parameters: inp_channels (`int`): Number of input channels. out_channels (`int`): Number of output channels. embed_dim (`int`): Embedding dimension. kernel_size (`int` or `tuple`): Size of the convolving kernel.
| 422 | |
| 423 | # unet_rl.py |
| 424 | class ResidualTemporalBlock1D(nn.Module): |
| 425 | """ |
| 426 | Residual 1D block with temporal convolutions. |
| 427 | |
| 428 | Parameters: |
| 429 | inp_channels (`int`): Number of input channels. |
| 430 | out_channels (`int`): Number of output channels. |
| 431 | embed_dim (`int`): Embedding dimension. |
| 432 | kernel_size (`int` or `tuple`): Size of the convolving kernel. |
| 433 | activation (`str`, defaults `mish`): It is possible to choose the right activation function. |
| 434 | """ |
| 435 | |
| 436 | def __init__( |
| 437 | self, |
| 438 | inp_channels: int, |
| 439 | out_channels: int, |
| 440 | embed_dim: int, |
| 441 | kernel_size: Union[int, Tuple[int, int]] = 5, |
| 442 | activation: str = "mish", |
| 443 | ): |
| 444 | super().__init__() |
| 445 | self.conv_in = Conv1dBlock(inp_channels, out_channels, kernel_size) |
| 446 | self.conv_out = Conv1dBlock(out_channels, out_channels, kernel_size) |
| 447 | |
| 448 | self.time_emb_act = get_activation(activation) |
| 449 | self.time_emb = nn.Linear(embed_dim, out_channels) |
| 450 | |
| 451 | self.residual_conv = ( |
| 452 | nn.Conv1d(inp_channels, out_channels, 1) if inp_channels != out_channels else nn.Identity() |
| 453 | ) |
| 454 | |
| 455 | def forward(self, inputs: torch.Tensor, t: torch.Tensor) -> torch.Tensor: |
| 456 | """ |
| 457 | Args: |
| 458 | inputs : [ batch_size x inp_channels x horizon ] |
| 459 | t : [ batch_size x embed_dim ] |
| 460 | |
| 461 | returns: |
| 462 | out : [ batch_size x out_channels x horizon ] |
| 463 | """ |
| 464 | t = self.time_emb_act(t) |
| 465 | t = self.time_emb(t) |
| 466 | out = self.conv_in(inputs) + rearrange_dims(t) |
| 467 | out = self.conv_out(out) |
| 468 | return out + self.residual_conv(inputs) |
| 469 | |
| 470 | |
| 471 | class TemporalConvLayer(nn.Module): |