## U-Net
| 207 | |
| 208 | |
| 209 | class UNet(nn.Module): |
| 210 | """ |
| 211 | ## U-Net |
| 212 | """ |
| 213 | |
| 214 | def __init__(self, input_channels: int = 2, output_channels: int = 1, n_channels: int = 32, |
| 215 | ch_mults: Union[Tuple[int, ...], List[int]] = (1, 2, 2, 4), |
| 216 | n_blocks: int = 2, is_noise: bool = True): |
| 217 | """ |
| 218 | * `image_channels` is the number of channels in the image. $3$ for RGB. |
| 219 | * `n_channels` is number of channels in the initial feature map that we transform the image into |
| 220 | * `ch_mults` is the list of channel numbers at each resolution. The number of channels is `ch_mults[i] * n_channels` |
| 221 | * `is_attn` is a list of booleans that indicate whether to use attention at each resolution |
| 222 | * `n_blocks` is the number of `UpDownBlocks` at each resolution |
| 223 | """ |
| 224 | super().__init__() |
| 225 | |
| 226 | # Number of resolutions |
| 227 | n_resolutions = len(ch_mults) |
| 228 | |
| 229 | # Project image into feature map |
| 230 | self.image_proj = nn.Conv2d(input_channels, n_channels, kernel_size=(3, 3), padding=(1, 1)) |
| 231 | |
| 232 | # Time embedding layer. Time embedding has `n_channels * 4` channels |
| 233 | self.is_noise = is_noise |
| 234 | if is_noise: |
| 235 | self.time_emb = TimeEmbedding(n_channels * 4) |
| 236 | |
| 237 | # #### First half of U-Net - decreasing resolution |
| 238 | down = [] |
| 239 | # Number of channels |
| 240 | out_channels = in_channels = n_channels |
| 241 | # For each resolution |
| 242 | for i in range(n_resolutions): |
| 243 | # Number of output channels at this resolution |
| 244 | out_channels = n_channels * ch_mults[i] |
| 245 | # Add `n_blocks` |
| 246 | for _ in range(n_blocks): |
| 247 | down.append(DownBlock(in_channels, out_channels, n_channels * 4, is_noise=is_noise)) |
| 248 | in_channels = out_channels |
| 249 | # Down sample at all resolutions except the last |
| 250 | if i < n_resolutions - 1: |
| 251 | down.append(Downsample(in_channels)) |
| 252 | |
| 253 | # Combine the set of modules |
| 254 | self.down = nn.ModuleList(down) |
| 255 | |
| 256 | # Middle block |
| 257 | self.middle = MiddleBlock(out_channels, n_channels * 4, is_noise=False) |
| 258 | |
| 259 | # #### Second half of U-Net - increasing resolution |
| 260 | up = [] |
| 261 | # Number of channels |
| 262 | in_channels = out_channels |
| 263 | # For each resolution |
| 264 | for i in reversed(range(n_resolutions)): |
| 265 | # `n_blocks` at the same resolution |
| 266 | out_channels = n_channels * ch_mults[i] |