(
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,
give_pre_end=False,
tanh_out=False,
use_linear_attn=False,
attn_type="vanilla",
**ignorekwargs,
)
| 546 | |
| 547 | class Decoder(nn.Module): |
| 548 | def __init__( |
| 549 | self, |
| 550 | *, |
| 551 | ch, |
| 552 | out_ch, |
| 553 | ch_mult=(1, 2, 4, 8), |
| 554 | num_res_blocks, |
| 555 | attn_resolutions, |
| 556 | dropout=0.0, |
| 557 | resamp_with_conv=True, |
| 558 | in_channels, |
| 559 | resolution, |
| 560 | z_channels, |
| 561 | give_pre_end=False, |
| 562 | tanh_out=False, |
| 563 | use_linear_attn=False, |
| 564 | attn_type="vanilla", |
| 565 | **ignorekwargs, |
| 566 | ): |
| 567 | super().__init__() |
| 568 | if use_linear_attn: |
| 569 | attn_type = "linear" |
| 570 | self.ch = ch |
| 571 | self.temb_ch = 0 |
| 572 | self.num_resolutions = len(ch_mult) |
| 573 | self.num_res_blocks = num_res_blocks |
| 574 | self.resolution = resolution |
| 575 | self.in_channels = in_channels |
| 576 | self.give_pre_end = give_pre_end |
| 577 | self.tanh_out = tanh_out |
| 578 | |
| 579 | # compute in_ch_mult, block_in and curr_res at lowest res |
| 580 | in_ch_mult = (1,) + tuple(ch_mult) |
| 581 | block_in = ch * ch_mult[self.num_resolutions - 1] |
| 582 | curr_res = resolution // 2 ** (self.num_resolutions - 1) |
| 583 | self.z_shape = (1, z_channels, curr_res, curr_res) |
| 584 | print("Working with z of shape {} = {} dimensions.".format(self.z_shape, np.prod(self.z_shape))) |
| 585 | |
| 586 | make_attn_cls = self._make_attn() |
| 587 | make_resblock_cls = self._make_resblock() |
| 588 | make_conv_cls = self._make_conv() |
| 589 | # z to block_in |
| 590 | self.conv_in = torch.nn.Conv2d(z_channels, block_in, kernel_size=3, stride=1, padding=1) |
| 591 | |
| 592 | # middle |
| 593 | self.mid = nn.Module() |
| 594 | self.mid.block_1 = make_resblock_cls( |
| 595 | in_channels=block_in, |
| 596 | out_channels=block_in, |
| 597 | temb_channels=self.temb_ch, |
| 598 | dropout=dropout, |
| 599 | ) |
| 600 | self.mid.attn_1 = make_attn_cls(block_in, attn_type=attn_type) |
| 601 | self.mid.block_2 = make_resblock_cls( |
| 602 | in_channels=block_in, |
| 603 | out_channels=block_in, |
| 604 | temb_channels=self.temb_ch, |
| 605 | dropout=dropout, |
nothing calls this directly
no test coverage detected