class DP_GAN_Generator(nn.Module):
    def __init__(self, opt):
        super().__init__()
        self.opt = opt
        ch = opt.channels_G
        if self.opt.upsample_cx:
            self.channels = [16*ch, 16*ch, 16*ch, 8*ch, 4*ch, 2*ch, 1*ch]
            self.conv_cx = nn.Conv2d(ch, 16*ch, 3, 1, 1)
            self.norm_cx = nn.InstanceNorm2d(16*ch)
            self.actv_cx = nn.ReLU()
            self.up_cx = nn.Sequential(
                            ResidualBlock(16*ch, 16*ch),
                            ResidualBlock(16*ch, 16*ch),
                            ResidualBlock(16*ch, 16*ch),
                            ResidualBlock(16*ch, 16*ch)
                            )
        else:
            self.channels = [ch, ch, ch, ch//2, ch//2, ch//4, ch//8]
        self.conv_img = nn.Conv2d(self.channels[-1], 3, 3, padding=1)
        self.up = nn.Upsample(scale_factor=2)
        self.decoder = nn.ModuleList([])
        for i in range(len(self.channels)-1):
            self.decoder.append(ResnetBlock_with_SPADE(self.channels[i], self.channels[i+1], opt))   

        
        self.seg_pyrmid = nn.ModuleList([])
        if self.opt.add_semantic_map:#semantic map + masked_image
            self.encoder = encoder(4 ,ch, partial=True)
        else:
            self.encoder = encoder(3 ,ch, partial=True)
        
        self.bitEstimator = BitEstimator(channel=ch)
        
        #(b, 64 , 256, 256)
        for i in range(len(self.channels) - 1):
            self.seg_pyrmid.append(nn.Sequential(nn.Conv2d(ch, ch, 3, stride=1, padding=1), nn.BatchNorm2d(ch), nn.ReLU(inplace=True)))
            #维度减半

    def forward(self, input, label_map, image, z=None):
        seg = input
        if self.opt.gpu_ids != "-1":
            seg.cuda()

        if self.opt.add_semantic_map:
            mask_map = PP.generate_mask_map(image, self.opt.mask_ratio, self.opt.patch_size)
            if self.opt.gpu_ids != "-1":
                mask_map = torch.from_numpy(mask_map).cuda()
            masked_image = image*mask_map
            x = torch.cat((masked_image, label_map), dim=1)
        else:
            mask_map = PP.generate_mask_map(image, self.opt.mask_ratio, self.opt.patch_size)
            if self.opt.gpu_ids != "-1":
                mask_map = torch.from_numpy(mask_map).cuda()
            masked_image = image*mask_map
            x = masked_image
        
        # # 查看iamge patch label 
        # transform1 = transforms.ToPILImage()
        # for q in range(self.opt.batch_size):
        #     transform1(masked_image[q]).save('./masked_image/'+str(q)+'image.jpg')
        if self.opt.add_semantic_map:
            x = self.encoder(x, mask_map.repeat(1, 4, 1, 1))
        else:
            x = self.encoder(x, mask_map.repeat(1, 3, 1, 1))

        quant_noise_feature = torch.zeros(x.size()).cuda()
        quant_noise_feature = torch.nn.init.uniform_(torch.zeros_like(quant_noise_feature), -0.5, 0.5)

        if self.training:
            compressed_x = x + quant_noise_feature
        else:
            compressed_x = torch.round(x)

        pyrmid = []
        s = compressed_x
        for op in self.seg_pyrmid:
            s = op(s)
            pyrmid.append(s)
            s = self.up(s)

        #TODO:将compressed_x的channel增加再送到生成器
        if self.opt.upsample_cx:
            decode_x = compressed_x
            decode_x = self.actv_cx(self.norm_cx(self.conv_cx(decode_x)))
            for i, layer in enumerate(self.up_cx):
                decode_x = layer(decode_x)
        else:
            decode_x = compressed_x
        for i in range(self.opt.num_res_blocks):
            seg_f = torch.cat([F.interpolate(pyrmid[0], size=(decode_x.shape[2], decode_x.shape[3]), mode='bilinear'), pyrmid[i]], 1)#Conv(i and UP(0))
            # seg_f = pyrmid[i]

            decode_x = self.decoder[i](decode_x, seg_f)
            if i < self.opt.num_res_blocks-1:
                decode_x = self.up(decode_x)

        decode_x = self.conv_img(F.leaky_relu(decode_x, 2e-1))
        decode_x = F.tanh(decode_x)

        def iclr18_estimate_bits_z(z):
            prob = self.bitEstimator(z + 0.5) - self.bitEstimator(z - 0.5)
            total_bits = torch.sum(torch.clamp(-1.0 * torch.log(prob + 1e-10) / math.log(2.0), 0, 50))
            return total_bits, prob

        total_bits_feature, _ = iclr18_estimate_bits_z(compressed_x)
        im_shape = image.size()
        bpp_feature = total_bits_feature / (self.opt.batch_size * im_shape[2] * im_shape[3])
        return decode_x, bpp_feature
为什么没报错self.training