Tiny Autoencoder block used in [`AutoencoderTiny`]. It is a mini residual module consisting of plain conv + ReLU blocks. Args: in_channels (`int`): The number of input channels. out_channels (`int`): The number of output channels. act_fn (`str`): ` T
| 562 | |
| 563 | |
| 564 | class AutoencoderTinyBlock(nn.Module): |
| 565 | """ |
| 566 | Tiny Autoencoder block used in [`AutoencoderTiny`]. It is a mini residual module consisting of plain conv + ReLU |
| 567 | blocks. |
| 568 | |
| 569 | Args: |
| 570 | in_channels (`int`): The number of input channels. |
| 571 | out_channels (`int`): The number of output channels. |
| 572 | act_fn (`str`): |
| 573 | ` The activation function to use. Supported values are `"swish"`, `"mish"`, `"gelu"`, and `"relu"`. |
| 574 | |
| 575 | Returns: |
| 576 | `torch.FloatTensor`: A tensor with the same shape as the input tensor, but with the number of channels equal to |
| 577 | `out_channels`. |
| 578 | """ |
| 579 | |
| 580 | def __init__(self, in_channels: int, out_channels: int, act_fn: str): |
| 581 | super().__init__() |
| 582 | act_fn = get_activation(act_fn) |
| 583 | self.conv = nn.Sequential( |
| 584 | nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1), |
| 585 | act_fn, |
| 586 | nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1), |
| 587 | act_fn, |
| 588 | nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1), |
| 589 | ) |
| 590 | self.skip = ( |
| 591 | nn.Conv2d(in_channels, out_channels, kernel_size=1, bias=False) |
| 592 | if in_channels != out_channels |
| 593 | else nn.Identity() |
| 594 | ) |
| 595 | self.fuse = nn.ReLU() |
| 596 | |
| 597 | def forward(self, x: torch.FloatTensor) -> torch.FloatTensor: |
| 598 | return self.fuse(self.conv(x) + self.skip(x)) |
| 599 | |
| 600 | |
| 601 | class UNetMidBlock2D(nn.Module): |