| 466 | |
| 467 | |
| 468 | class Decoder(nn.Module): |
| 469 | def __init__(self, *, ch, out_ch, ch_mult=(1,2,4,8), num_res_blocks, |
| 470 | attn_resolutions, dropout=0.0, resamp_with_conv=True, in_channels, |
| 471 | resolution, z_channels, give_pre_end=False, tanh_out=False, use_linear_attn=False, |
| 472 | attn_type="vanilla", **ignorekwargs): |
| 473 | super().__init__() |
| 474 | if use_linear_attn: attn_type = "linear" |
| 475 | self.ch = ch |
| 476 | self.temb_ch = 0 |
| 477 | self.num_resolutions = len(ch_mult) |
| 478 | self.num_res_blocks = num_res_blocks |
| 479 | self.resolution = resolution |
| 480 | self.in_channels = in_channels |
| 481 | self.give_pre_end = give_pre_end |
| 482 | self.tanh_out = tanh_out |
| 483 | |
| 484 | # compute in_ch_mult, block_in and curr_res at lowest res |
| 485 | in_ch_mult = (1,)+tuple(ch_mult) |
| 486 | block_in = ch*ch_mult[self.num_resolutions-1] |
| 487 | curr_res = resolution // 2**(self.num_resolutions-1) |
| 488 | self.z_shape = (1,z_channels,curr_res,curr_res) |
| 489 | print("AE working on z of shape {} = {} dimensions.".format( |
| 490 | self.z_shape, np.prod(self.z_shape))) |
| 491 | |
| 492 | # z to block_in |
| 493 | self.conv_in = torch.nn.Conv2d(z_channels, |
| 494 | block_in, |
| 495 | kernel_size=3, |
| 496 | stride=1, |
| 497 | padding=1) |
| 498 | |
| 499 | # middle |
| 500 | self.mid = nn.Module() |
| 501 | self.mid.block_1 = ResnetBlock(in_channels=block_in, |
| 502 | out_channels=block_in, |
| 503 | temb_channels=self.temb_ch, |
| 504 | dropout=dropout) |
| 505 | self.mid.attn_1 = make_attn(block_in, attn_type=attn_type) |
| 506 | self.mid.block_2 = ResnetBlock(in_channels=block_in, |
| 507 | out_channels=block_in, |
| 508 | temb_channels=self.temb_ch, |
| 509 | dropout=dropout) |
| 510 | |
| 511 | # upsampling |
| 512 | self.up = nn.ModuleList() |
| 513 | for i_level in reversed(range(self.num_resolutions)): |
| 514 | block = nn.ModuleList() |
| 515 | attn = nn.ModuleList() |
| 516 | block_out = ch*ch_mult[i_level] |
| 517 | for i_block in range(self.num_res_blocks+1): |
| 518 | block.append(ResnetBlock(in_channels=block_in, |
| 519 | out_channels=block_out, |
| 520 | temb_channels=self.temb_ch, |
| 521 | dropout=dropout)) |
| 522 | block_in = block_out |
| 523 | if curr_res in attn_resolutions: |
| 524 | attn.append(make_attn(block_in, attn_type=attn_type)) |
| 525 | up = nn.Module() |