(self, *, ch, out_ch, ch_mult=(1,2,4,8), num_res_blocks,
attn_resolutions, dropout=0.0, resamp_with_conv=True, in_channels,
resolution, z_channels, double_z=True, use_linear_attn=False, attn_type="vanilla",
**ignore_kwargs)
| 365 | |
| 366 | class Encoder(nn.Module): |
| 367 | def __init__(self, *, ch, out_ch, ch_mult=(1,2,4,8), num_res_blocks, |
| 368 | attn_resolutions, dropout=0.0, resamp_with_conv=True, in_channels, |
| 369 | resolution, z_channels, double_z=True, use_linear_attn=False, attn_type="vanilla", |
| 370 | **ignore_kwargs): |
| 371 | super().__init__() |
| 372 | if use_linear_attn: attn_type = "linear" |
| 373 | self.ch = ch |
| 374 | self.temb_ch = 0 |
| 375 | self.num_resolutions = len(ch_mult) |
| 376 | self.num_res_blocks = num_res_blocks |
| 377 | self.resolution = resolution |
| 378 | self.in_channels = in_channels |
| 379 | |
| 380 | # downsampling |
| 381 | self.conv_in = torch.nn.Conv2d(in_channels, |
| 382 | self.ch, |
| 383 | kernel_size=3, |
| 384 | stride=1, |
| 385 | padding=1) |
| 386 | |
| 387 | curr_res = resolution |
| 388 | in_ch_mult = (1,)+tuple(ch_mult) |
| 389 | self.in_ch_mult = in_ch_mult |
| 390 | self.down = nn.ModuleList() |
| 391 | for i_level in range(self.num_resolutions): |
| 392 | block = nn.ModuleList() |
| 393 | attn = nn.ModuleList() |
| 394 | block_in = ch*in_ch_mult[i_level] |
| 395 | block_out = ch*ch_mult[i_level] |
| 396 | for i_block in range(self.num_res_blocks): |
| 397 | block.append(ResnetBlock(in_channels=block_in, |
| 398 | out_channels=block_out, |
| 399 | temb_channels=self.temb_ch, |
| 400 | dropout=dropout)) |
| 401 | block_in = block_out |
| 402 | if curr_res in attn_resolutions: |
| 403 | attn.append(make_attn(block_in, attn_type=attn_type)) |
| 404 | down = nn.Module() |
| 405 | down.block = block |
| 406 | down.attn = attn |
| 407 | if i_level != self.num_resolutions-1: |
| 408 | down.downsample = Downsample(block_in, resamp_with_conv) |
| 409 | curr_res = curr_res // 2 |
| 410 | self.down.append(down) |
| 411 | |
| 412 | # middle |
| 413 | self.mid = nn.Module() |
| 414 | self.mid.block_1 = ResnetBlock(in_channels=block_in, |
| 415 | out_channels=block_in, |
| 416 | temb_channels=self.temb_ch, |
| 417 | dropout=dropout) |
| 418 | self.mid.attn_1 = make_attn(block_in, attn_type=attn_type) |
| 419 | self.mid.block_2 = ResnetBlock(in_channels=block_in, |
| 420 | out_channels=block_in, |
| 421 | temb_channels=self.temb_ch, |
| 422 | dropout=dropout) |
| 423 | |
| 424 | # end |
nothing calls this directly
no test coverage detected