示例#1
0
文件: net.py 项目: Cristo-R/cnn
    def __init__(self, in_channels, act=F.elu, norm=None, opt=None):
        super(InpaintingModel_GMCNN, self).__init__()
        self.opt = opt
        self.init(opt)

        self.confidence_mask_layer = ConfidenceDrivenMaskLayer()

        self.netGM = GMCNN(in_channels,
                           out_channels=3,
                           cnum=opt.g_cnum,
                           act=act,
                           norm=norm).cuda()
        init_weights(self.netGM)
        self.model_names = ['GM']
        if self.opt.phase == 'test':
            return

        self.netD = None

        self.optimizer_G = torch.optim.Adam(self.netGM.parameters(),
                                            lr=opt.lr,
                                            betas=(0.5, 0.9))
        self.optimizer_D = None

        self.wganloss = None
        self.recloss = nn.L1Loss()
        self.aeloss = nn.L1Loss()
        self.mrfloss = None
        self.lambda_adv = opt.lambda_adv
        self.lambda_rec = opt.lambda_rec
        self.lambda_ae = opt.lambda_ae
        self.lambda_gp = opt.lambda_gp
        self.lambda_mrf = opt.lambda_mrf
        self.G_loss = None
        self.G_loss_reconstruction = None
        self.G_loss_mrf = None
        self.G_loss_adv, self.G_loss_adv_local = None, None
        self.G_loss_ae = None
        self.D_loss, self.D_loss_local = None, None
        self.GAN_loss = None

        self.gt, self.gt_local = None, None
        self.mask, self.mask_01 = None, None
        self.rect = None
        self.im_in, self.gin = None, None

        self.completed, self.completed_local = None, None
        self.completed_logit, self.completed_local_logit = None, None
        self.gt_logit, self.gt_local_logit = None, None

        self.pred = None

        if self.opt.pretrain_network is False:
            if self.opt.mask_type == 'rect':
                self.netD = GlobalLocalDiscriminator(
                    3,
                    cnum=opt.d_cnum,
                    act=act,
                    g_fc_channels=opt.img_shapes[0] // 16 *
                    opt.img_shapes[1] // 16 * opt.d_cnum * 4,
                    l_fc_channels=opt.mask_shapes[0] // 16 *
                    opt.mask_shapes[1] // 16 * opt.d_cnum * 4,
                    spectral_norm=self.opt.spectral_norm).cuda()
            else:
                self.netD = GlobalLocalDiscriminator(
                    3,
                    cnum=opt.d_cnum,
                    act=act,
                    spectral_norm=self.opt.spectral_norm,
                    g_fc_channels=opt.img_shapes[0] // 16 *
                    opt.img_shapes[1] // 16 * opt.d_cnum * 4,
                    l_fc_channels=opt.img_shapes[0] // 16 *
                    opt.img_shapes[1] // 16 * opt.d_cnum * 4).cuda()
            init_weights(self.netD)
            self.optimizer_D = torch.optim.Adam(filter(
                lambda x: x.requires_grad, self.netD.parameters()),
                                                lr=opt.lr,
                                                betas=(0.5, 0.9))
            self.wganloss = WGANLoss()
            self.mrfloss = IDMRFLoss()
示例#2
0
文件: net.py 项目: yhl41001/DMFN
    def __init__(self, act=F.elu, opt=None):
        super(InpaintingModel_DFBM, self).__init__()
        self.opt = opt
        self.init(opt)

        self.confidence_mask_layer = ConfidenceDrivenMaskLayer()

        self.netDFBN = DFBN().cuda()
        init_weights(self.netDFBN)
        self.model_names = ['DFBN']
        if self.opt.phase == 'test':
            return

        self.netD = None
        self.optimizers = []
        self.optimizer_G = torch.optim.Adam(self.netDFBN.parameters(),
                                            lr=opt.lr,
                                            betas=(0.5, 0.9))
        self.optimizers += [self.optimizer_G]
        self.optimizer_D = None
        self.zeros = torch.zeros((opt.batch_size, 1)).cuda()
        self.ones = torch.ones((opt.batch_size, 1)).cuda()
        self.aeloss = nn.L1Loss()
        self.vggloss = None
        self.G_loss = None
        self.G_loss_mrf = None
        self.G_loss_adv, self.G_loss_vgg, self.G_loss_fm_dis = None, None, None
        self.G_loss_ae = None
        self.loss_eta = 5
        self.loss_mu = 0.03
        self.loss_vgg = 1
        self.BCEloss = nn.BCEWithLogitsLoss().cuda()
        self.gt, self.gt_local = None, None
        self.mask, self.mask_01 = None, None
        self.rect = None
        self.im_in, self.gin = None, None

        self.completed, self.completed_local = None, None
        self.completed_logit, self.gt_logit = None, None

        def weight_fn(layer):
            s = layer.shape
            return 1e3 / (s[1] * s[1] * s[1] * s[2] * s[3])

        self.weight_fn = weight_fn

        self.pred = None

        self.netD = GlobalLocalDiscriminator(3,
                                             cnum=opt.d_cnum,
                                             act=F.leaky_relu).cuda()
        init_weights(self.netD)
        self.optimizer_D = torch.optim.Adam(filter(lambda x: x.requires_grad,
                                                   self.netD.parameters()),
                                            lr=opt.lr,
                                            betas=(0.5, 0.9))
        self.vggloss = VGGLoss()
        self.optimizers += [self.optimizer_D]
        self.schedulers = []
        for optimizer in self.optimizers:
            self.schedulers.append(
                lr_scheduler.MultiStepLR(optimizer, [2000, 40000], 0.5))