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
| 550 | |
| 551 | |
| 552 | class AutoencoderTinyBlock(nn.Module): |
| 553 | """ |
| 554 | Tiny Autoencoder block used in [`AutoencoderTiny`]. It is a mini residual module consisting of plain conv + ReLU |
| 555 | blocks. |
| 556 | |
| 557 | Args: |
| 558 | in_channels (`int`): The number of input channels. |
| 559 | out_channels (`int`): The number of output channels. |
| 560 | act_fn (`str`): |
| 561 | ` The activation function to use. Supported values are `"swish"`, `"mish"`, `"gelu"`, and `"relu"`. |
| 562 | |
| 563 | Returns: |
| 564 | `torch.Tensor`: A tensor with the same shape as the input tensor, but with the number of channels equal to |
| 565 | `out_channels`. |
| 566 | """ |
| 567 | |
| 568 | def __init__(self, in_channels: int, out_channels: int, act_fn: str): |
| 569 | super().__init__() |
| 570 | act_fn = get_activation(act_fn) |
| 571 | self.conv = nn.Sequential( |
| 572 | nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1), |
| 573 | act_fn, |
| 574 | nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1), |
| 575 | act_fn, |
| 576 | nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1), |
| 577 | ) |
| 578 | self.skip = ( |
| 579 | nn.Conv2d(in_channels, out_channels, kernel_size=1, bias=False) |
| 580 | if in_channels != out_channels |
| 581 | else nn.Identity() |
| 582 | ) |
| 583 | self.fuse = nn.ReLU() |
| 584 | |
| 585 | def forward(self, x: torch.Tensor) -> torch.Tensor: |
| 586 | return self.fuse(self.conv(x) + self.skip(x)) |
| 587 | |
| 588 | |
| 589 | class UNetMidBlock2D(nn.Module): |