(
self,
in_channels,
out_channels,
ch,
num_res_blocks,
resolution,
ch_mult=(2, 2),
dropout=0.0,
)
| 749 | |
| 750 | class UpsampleDecoder(nn.Module): |
| 751 | def __init__( |
| 752 | self, |
| 753 | in_channels, |
| 754 | out_channels, |
| 755 | ch, |
| 756 | num_res_blocks, |
| 757 | resolution, |
| 758 | ch_mult=(2, 2), |
| 759 | dropout=0.0, |
| 760 | ): |
| 761 | super().__init__() |
| 762 | # upsampling |
| 763 | self.temb_ch = 0 |
| 764 | self.num_resolutions = len(ch_mult) |
| 765 | self.num_res_blocks = num_res_blocks |
| 766 | block_in = in_channels |
| 767 | curr_res = resolution // 2 ** (self.num_resolutions - 1) |
| 768 | self.res_blocks = nn.ModuleList() |
| 769 | self.upsample_blocks = nn.ModuleList() |
| 770 | for i_level in range(self.num_resolutions): |
| 771 | res_block = [] |
| 772 | block_out = ch * ch_mult[i_level] |
| 773 | for i_block in range(self.num_res_blocks + 1): |
| 774 | res_block.append( |
| 775 | ResnetBlock( |
| 776 | in_channels=block_in, |
| 777 | out_channels=block_out, |
| 778 | temb_channels=self.temb_ch, |
| 779 | dropout=dropout, |
| 780 | ) |
| 781 | ) |
| 782 | block_in = block_out |
| 783 | self.res_blocks.append(nn.ModuleList(res_block)) |
| 784 | if i_level != self.num_resolutions - 1: |
| 785 | self.upsample_blocks.append(Upsample(block_in, True)) |
| 786 | curr_res = curr_res * 2 |
| 787 | |
| 788 | # end |
| 789 | self.norm_out = Normalize(block_in) |
| 790 | self.conv_out = torch.nn.Conv2d( |
| 791 | block_in, out_channels, kernel_size=3, stride=1, padding=1 |
| 792 | ) |
| 793 | |
| 794 | def forward(self, x): |
| 795 | # upsampling |
nothing calls this directly
no test coverage detected