(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, output_key=[],
skipconnect_type=None, skipconnect_emb=None, skipconnect_layer=None, use_skipconnect_proj=False, **ignorekwargs)
| 917 | |
| 918 | class UNetDecoder(nn.Module): |
| 919 | def __init__(self, *, ch, out_ch, ch_mult=(1, 2, 4, 8), num_res_blocks, |
| 920 | attn_resolutions, dropout=0.0, resamp_with_conv=True, in_channels, |
| 921 | resolution, z_channels, give_pre_end=False, output_key=[], |
| 922 | skipconnect_type=None, skipconnect_emb=None, skipconnect_layer=None, use_skipconnect_proj=False, **ignorekwargs): |
| 923 | super().__init__() |
| 924 | self.output_key = output_key |
| 925 | self.ch = ch |
| 926 | self.temb_ch = 0 |
| 927 | self.num_resolutions = len(ch_mult) |
| 928 | self.num_res_blocks = num_res_blocks |
| 929 | self.resolution = resolution |
| 930 | self.in_channels = in_channels |
| 931 | self.give_pre_end = give_pre_end |
| 932 | self.skipconnect_type = skipconnect_type # {sum, concat, nonlinear_sum} |
| 933 | self.skipconnect_emb = skipconnect_emb |
| 934 | self.skipconnect_layer = skipconnect_layer |
| 935 | self.use_skipconnect_proj = use_skipconnect_proj |
| 936 | assert self.skipconnect_type in ['sum', 'concat'] |
| 937 | |
| 938 | # compute in_ch_mult, block_in and curr_res at lowest res |
| 939 | in_ch_mult = (1,)+tuple(ch_mult) |
| 940 | block_in = ch*ch_mult[self.num_resolutions-1] |
| 941 | curr_res = resolution // 2**(self.num_resolutions-1) |
| 942 | self.z_shape = (1, z_channels, curr_res, curr_res) |
| 943 | print("Working with z of shape {} = {} dimensions.".format( |
| 944 | self.z_shape, np.prod(self.z_shape))) |
| 945 | |
| 946 | self.conv_in = torch.nn.Conv2d(z_channels, |
| 947 | block_in, |
| 948 | kernel_size=3, |
| 949 | stride=1, |
| 950 | padding=1) |
| 951 | |
| 952 | # middle |
| 953 | self.mid = nn.Module() |
| 954 | self.mid.block_1 = ResnetBlock(in_channels=block_in, |
| 955 | out_channels=block_in, |
| 956 | temb_channels=self.temb_ch, |
| 957 | dropout=dropout) |
| 958 | self.mid.attn_1 = AttnBlock(block_in) |
| 959 | self.mid.block_2 = ResnetBlock(in_channels=block_in, |
| 960 | out_channels=block_in, |
| 961 | temb_channels=self.temb_ch, |
| 962 | dropout=dropout) |
| 963 | |
| 964 | # upsampling |
| 965 | skip_ch = 0 |
| 966 | self.up = nn.ModuleList() |
| 967 | self.skip_proj = nn.ModuleDict() |
| 968 | for i, i_level in enumerate(reversed(range(self.num_resolutions))): |
| 969 | proj = [] |
| 970 | block = nn.ModuleList() |
| 971 | attn = nn.ModuleList() |
| 972 | block_out = ch*ch_mult[i_level] |
| 973 | if i != 0: |
| 974 | skip_ch = skipconnect_emb[i-1] |
| 975 | if (self.skipconnect_type == 'sum') & (skip_ch != block_in) & (not self.use_skipconnect_proj): |
| 976 | self.skip_proj[str(i_level)] = nn.Conv2d(in_channels=skip_ch, |
nothing calls this directly
no test coverage detected