* `in_channels` is the number of input channels * `out_channels` is the number of input channels * `time_channels` is the number channels in the time step ($t$) embeddings * `n_groups` is the number of groups for [group normalization](../../normalization/group_no
(self, in_channels: int, out_channels: int, time_channels: int,
dropout: float = 0.1, is_noise: bool = True)
| 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 | """ |