### Residual block A residual block has two convolution layers with group normalization. Each resolution is processed with two residual blocks.
| 59 | |
| 60 | |
| 61 | class ResidualBlock(nn.Module): |
| 62 | """ |
| 63 | ### Residual block |
| 64 | A residual block has two convolution layers with group normalization. |
| 65 | Each resolution is processed with two residual blocks. |
| 66 | """ |
| 67 | |
| 68 | def __init__(self, in_channels: int, out_channels: int, time_channels: int, |
| 69 | dropout: float = 0.1, is_noise: bool = True): |
| 70 | """ |
| 71 | * `in_channels` is the number of input channels |
| 72 | * `out_channels` is the number of input channels |
| 73 | * `time_channels` is the number channels in the time step ($t$) embeddings |
| 74 | * `n_groups` is the number of groups for [group normalization](../../normalization/group_norm/index.html) |
| 75 | * `dropout` is the dropout rate |
| 76 | """ |
| 77 | super().__init__() |
| 78 | # Group normalization and the first convolution layer |
| 79 | self.is_noise = is_noise |
| 80 | self.act1 = Swish() |
| 81 | self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=(3, 3), padding=(1, 1)) |
| 82 | |
| 83 | # Group normalization and the second convolution layer |
| 84 | |
| 85 | self.act2 = Swish() |
| 86 | self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=(3, 3), padding=(1, 1)) |
| 87 | |
| 88 | # If the number of input channels is not equal to the number of output channels we have to |
| 89 | # project the shortcut connection |
| 90 | if in_channels != out_channels: |
| 91 | self.shortcut = nn.Conv2d(in_channels, out_channels, kernel_size=(1, 1)) |
| 92 | else: |
| 93 | self.shortcut = nn.Identity() |
| 94 | |
| 95 | # Linear layer for time embeddings |
| 96 | if self.is_noise: |
| 97 | self.time_emb = nn.Linear(time_channels, out_channels) |
| 98 | self.time_act = Swish() |
| 99 | |
| 100 | self.dropout = nn.Dropout(dropout) |
| 101 | |
| 102 | def forward(self, x: torch.Tensor, t: torch.Tensor): |
| 103 | """ |
| 104 | * `x` has shape `[batch_size, in_channels, height, width]` |
| 105 | * `t` has shape `[batch_size, time_channels]` |
| 106 | """ |
| 107 | # First convolution layer |
| 108 | h = self.conv1(self.act1(x)) |
| 109 | # Add time embeddings |
| 110 | if self.is_noise: |
| 111 | h += self.time_emb(self.time_act(t))[:, :, None, None] |
| 112 | # Second convolution layer |
| 113 | h = self.conv2(self.dropout(self.act2(h))) |
| 114 | |
| 115 | # Add the shortcut connection and return |
| 116 | return h + self.shortcut(x) |
| 117 | |
| 118 |