<commit_before>import numpy as np
import torch
import os
from collections import OrderedDict
from torch.autograd import Variable
import util.util as util
from util.image_pool import ImagePool
from .base_model import BaseModel
from . import networks


class Pix2PixModel(BaseModel):
    def name(self):
        return 'Pix2PixModel'

    def initialize(self, opt):
        BaseModel.initialize(self, opt)
        self.isTrain = opt.isTrain
        # define tensors
        self.input_A = self.Tensor(opt.batchSize, opt.input_nc,
                                   opt.fineSize, opt.fineSize)
        self.input_B = self.Tensor(opt.batchSize, opt.output_nc,
                                   opt.fineSize, opt.fineSize)
        self.loss_G_GAN = Variable(torch.FloatTensor([0]).cuda())
        self.loss_G_L1 = Variable(torch.FloatTensor([0]).cuda())
        # load/define networks
        self.netG = networks.define_G(3, opt.output_nc, opt.ngf,
                                      opt.which_model_netG, opt.norm, not opt.no_dropout, opt.init_type, self.gpu_ids)
        if self.isTrain:
            # use_sigmoid = opt.no_lsgan
            use_sigmoid = False
            self.netD = networks.define_D(opt.input_nc + opt.output_nc, opt.ndf,
                                          opt.which_model_netD,
                                          opt.n_layers_D, opt.norm, use_sigmoid, opt.init_type, self.gpu_ids)

        if not self.isTrain or opt.continue_train:
            self.load_network(self.netG, 'G', opt.which_epoch)
            if self.isTrain:
                self.load_network(self.netD, 'D', opt.which_epoch)

        if self.isTrain:
            self.fake_AB_pool = ImagePool(opt.pool_size)
            self.old_lr = opt.lr
            # define loss functions
            self.criterionGAN = networks.GANLoss(use_lsgan=not opt.no_lsgan, tensor=self.Tensor)
            self.criterionL1 = torch.nn.L1Loss()

            # initialize optimizers
            self.schedulers = []
            self.optimizers = []
            # self.optimizer_G = torch.optim.Adam(self.netG.parameters(),
            #                                     lr=opt.lr, betas=(opt.beta1, 0.999))
            # self.optimizer_D = torch.optim.Adam(self.netD.parameters(),
            #                                     lr=opt.lr, betas=(opt.beta1, 0.999))
            self.optimizer_G = torch.optim.RMSprop(self.netG.parameters(), lr=opt.lr)
            self.optimizer_D = torch.optim.RMSprop(self.netD.parameters(), lr=opt.lr)
            self.optimizers.append(self.optimizer_G)
            self.optimizers.append(self.optimizer_D)
            # for optimizer in self.optimizers:
            #     self.schedulers.append(networks.get_scheduler(optimizer, opt))

        print('---------- Networks initialized -------------')
        networks.print_network(self.netG)
        if self.isTrain:
            networks.print_network(self.netD)
        print('-----------------------------------------------')

    def set_input(self, input):
        AtoB = self.opt.which_direction == 'AtoB'
        input_A = input['A' if AtoB else 'B']
        input_B = input['B' if AtoB else 'A']
        self.input_A.resize_(input_A.size()).copy_(input_A)
        self.input_B.resize_(input_B.size()).copy_(input_B)
        self.image_paths = input['A_paths' if AtoB else 'B_paths']

    def forward(self):
        self.real_A = Variable(self.input_A)
        self.fake_B = self.netG.forward(self.real_A)
        self.real_B = Variable(self.input_B)

    def forward_with_noise(self):
        self.real_A = Variable(self.input_A)
        batch_size = self.real_A.size(0)
        fixed_noise = Variable(torch.FloatTensor(batch_size, 3, 256, 256).normal_(0, 1).cuda())
        real_A_with_noise = torch.cat((self.real_A, fixed_noise), 1)
        self.fake_B = self.netG.forward(real_A_with_noise)
        self.real_B = Variable(self.input_B)

    # no backprop gradients
    def test(self):
        self.real_A = Variable(self.input_A, volatile=True)
        self.fake_B = self.netG.forward(self.real_A)
        self.real_B = Variable(self.input_B, volatile=True)

    # get image paths
    def get_image_paths(self):
        return self.image_paths

    def backward_D(self):
        # Fake
        # stop backprop to the generator by detaching fake_B
        fake_AB = self.fake_AB_pool.query(torch.cat((self.real_A, self.fake_B), 1))
        self.pred_fake = self.netD.forward(fake_AB.detach())
        self.loss_D_fake = self.criterionGAN(self.pred_fake, False)
        # Real
        real_AB = torch.cat((self.real_A, self.real_B), 1)
        self.pred_real = self.netD.forward(real_AB)
        self.loss_D_real = self.criterionGAN(self.pred_real, True)

        # Combined loss
        self.loss_D = (self.loss_D_fake + self.loss_D_real) * 0.5

        self.loss_D.backward()

    def backward_wgan_D(self):
        real_AB = torch.cat((self.real_A, self.real_B), 1)
        self.loss_D_real = self.netD.forward(real_AB)
        # self.loss_D_real = self.loss_D_real.view(-1, 1).mean(0)
        # The Numbers 1
        one = torch.FloatTensor([1]).cuda()
        self.loss_D_real.backward(one)
        # Fake
        # stop backprop to the generator by detaching fake_B
        fake_AB = self.fake_AB_pool.query(torch.cat((self.real_A, self.fake_B), 1))
        self.loss_D_fake = self.netD.forward(fake_AB.detach())
        # self.loss_D_fake = self.loss_D_fake.view(-1, 1).mean(0)
        # The numbers -1
        mone = one * -1
        self.loss_D_fake.backward(mone)
        self.loss_D = self.loss_D_real - self.loss_D_fake

    def backward_G(self):
        # First, G(A) should fake the discriminator
        fake_AB = torch.cat((self.real_A, self.fake_B), 1)
        pred_fake = self.netD.forward(fake_AB)
        self.loss_G_GAN = self.criterionGAN(pred_fake, True)

        # Second, G(A) = B
        self.loss_G_L1 = self.criterionL1(self.fake_B, self.real_B) * self.opt.lambda_A

        self.loss_G = self.loss_G_GAN + self.loss_G_L1

        self.loss_G.backward()

    def backward_wgan_G(self):
        # First, G(A) should fake the discriminator
        fake_AB = torch.cat((self.real_A, self.fake_B), 1)
        self.loss_G_GAN = self.netD.forward(fake_AB)
        # self.loss_G_GAN = self.loss_G_GAN.view(-1, 1).mean(0)
        # self.loss_G_L1 = self.criterionL1(self.fake_B, self.real_B) * self.opt.lambda_A
        # self.loss_G = self.loss_G_GAN + self.loss_G_L1
        self.loss_G = self.loss_G_GAN
        # The Numbers 1
        one = torch.FloatTensor([1]).cuda()
        self.loss_G.backward(one)

    def optimize_parameters(self):
        self.forward()

        self.optimizer_D.zero_grad()
        self.backward_D()
        self.optimizer_D.step()

        self.optimizer_G.zero_grad()
        self.backward_G()
        self.optimizer_G.step()

    def optimize_netD_parameters(self):
        make_trainable(self.netD, True)
        for p in self.netD.parameters():
            p.data.clamp(-0.01, 0.01)
        self.forward()
        self.optimizer_D.zero_grad()
        self.backward_wgan_D()<commit_msg>debug param setting for improved pix2pix<commit_after>import numpy as np
import torch
import os
from collections import OrderedDict
from torch.autograd import Variable
import util.util as util
from util.image_pool import ImagePool
from .base_model import BaseModel
from . import networks


class Pix2PixModel(BaseModel):
    def name(self):
        return 'Pix2PixModel'

    def initialize(self, opt):
        BaseModel.initialize(self, opt)
        self.isTrain = opt.isTrain
        # define tensors
        self.input_A = self.Tensor(opt.batchSize, opt.input_nc,
                                   opt.fineSize, opt.fineSize)
        self.input_B = self.Tensor(opt.batchSize, opt.output_nc,
                                   opt.fineSize, opt.fineSize)
        self.loss_G_GAN = Variable(torch.FloatTensor([0]).cuda())
        self.loss_G_L1 = Variable(torch.FloatTensor([0]).cuda())
        # load/define networks
        self.netG = networks.define_G(3, opt.output_nc, opt.ngf,
                                      opt.which_model_netG, opt.norm, not opt.no_dropout, opt.init_type, self.gpu_ids)
        if self.isTrain:
            # use_sigmoid = opt.no_lsgan
            use_sigmoid = False
            self.netD = networks.define_D(opt.input_nc + opt.output_nc, opt.ndf,
                                          opt.which_model_netD,
                                          opt.n_layers_D, opt.norm, use_sigmoid, opt.init_type, self.gpu_ids)

        if not self.isTrain or opt.continue_train:
            self.load_network(self.netG, 'G', opt.which_epoch)
            if self.isTrain:
                self.load_network(self.netD, 'D', opt.which_epoch)

        if self.isTrain:
            self.fake_AB_pool = ImagePool(opt.pool_size)
            self.old_lr = opt.lr
            # define loss functions
            self.criterionGAN = networks.GANLoss(use_lsgan=not opt.no_lsgan, tensor=self.Tensor)
            self.criterionL1 = torch.nn.L1Loss()

            # initialize optimizers
            self.schedulers = []
            self.optimizers = []
            # self.optimizer_G = torch.optim.Adam(self.netG.parameters(),
            #                                     lr=opt.lr, betas=(opt.beta1, 0.999))
            # self.optimizer_D = torch.optim.Adam(self.netD.parameters(),
            #                                     lr=opt.lr, betas=(opt.beta1, 0.999))
            self.optimizer_G = torch.optim.RMSprop(self.netG.parameters(), lr=opt.lr)
            self.optimizer_D = torch.optim.RMSprop(self.netD.parameters(), lr=opt.lr)
            self.optimizers.append(self.optimizer_G)
            self.optimizers.append(self.optimizer_D)
            # for optimizer in self.optimizers:
            #     self.schedulers.append(networks.get_scheduler(optimizer, opt))

        print('---------- Networks initialized -------------')
        networks.print_network(self.netG)
        if self.isTrain:
            networks.print_network(self.netD)
        print('-----------------------------------------------')

    def set_input(self, input):
        AtoB = self.opt.which_direction == 'AtoB'
        input_A = input['A' if AtoB else 'B']
        input_B = input['B' if AtoB else 'A']
        self.input_A.resize_(input_A.size()).copy_(input_A)
        self.input_B.resize_(input_B.size()).copy_(input_B)
        self.image_paths = input['A_paths' if AtoB else 'B_paths']

    def forward(self):
        self.real_A = Variable(self.input_A)
        self.fake_B = self.netG.forward(self.real_A)
        self.real_B = Variable(self.input_B)

    def forward_with_noise(self):
        self.real_A = Variable(self.input_A)
        batch_size = self.real_A.size(0)
        fixed_noise = Variable(torch.FloatTensor(batch_size, 3, 256, 256).normal_(0, 1).cuda())
        real_A_with_noise = torch.cat((self.real_A, fixed_noise), 1)
        self.fake_B = self.netG.forward(real_A_with_noise)
        self.real_B = Variable(self.input_B)

    # no backprop gradients
    def test(self):
        self.real_A = Variable(self.input_A, volatile=True)
        self.fake_B = self.netG.forward(self.real_A)
        self.real_B = Variable(self.input_B, volatile=True)

    # get image paths
    def get_image_paths(self):
        return self.image_paths

    def backward_D(self):
        # Fake
        # stop backprop to the generator by detaching fake_B
        fake_AB = self.fake_AB_pool.query(torch.cat((self.real_A, self.fake_B), 1))
        self.pred_fake = self.netD.forward(fake_AB.detach())
        self.loss_D_fake = self.criterionGAN(self.pred_fake, False)
        # Real
        real_AB = torch.cat((self.real_A, self.real_B), 1)
        self.pred_real = self.netD.forward(real_AB)
        self.loss_D_real = self.criterionGAN(self.pred_real, True)

        # Combined loss
        self.loss_D = (self.loss_D_fake + self.loss_D_real) * 0.5

        self.loss_D.backward()

    def backward_wgan_D(self):
        real_AB = torch.cat((self.real_A, self.real_B), 1)
        self.loss_D_real = self.netD.forward(real_AB)
        # print('fffffffffffffffff', self.loss_D_real.size)
        self.loss_D_real = self.loss_D_real.view(-1, 1).mean(0)
        # The Numbers 1
        one = torch.FloatTensor([1]).cuda()
        self.loss_D_real.backward(one)
        # Fake
        # stop backprop to the generator by detaching fake_B
        fake_AB = self.fake_AB_pool.query(torch.cat((self.real_A, self.fake_B), 1))
        self.loss_D_fake = self.netD.forward(fake_AB.detach())
        self.loss_D_fake = self.loss_D_fake.view(-1, 1).mean(0)
        # The numbers -1
        mone = one * -1
        self.loss_D_fake.backward(mone)
        self.loss_D = self.loss_D_real - self.loss_D_fake

    def backward_G(self):
        # First, G(A) should fake the discriminator
        fake_AB = torch.cat((self.real_A, self.fake_B), 1)
        pred_fake = self.netD.forward(fake_AB)
        self.loss_G_GAN = self.criterionGAN(pred_fake, True)

        # Second, G(A) = B
        self.loss_G_L1 = self.criterionL1(self.fake_B, self.real_B) * self.opt.lambda_A

        self.loss_G = self.loss_G_GAN + self.loss_G_L1

        self.loss_G.backward()

    def backward_wgan_G(self):
        # First, G(A) should fake the discriminator
        fake_AB = torch.cat((self.real_A, self.fake_B), 1)
        self.loss_G_GAN = self.netD.forward(fake_AB)
        self.loss_G_GAN = self.loss_G_GAN.view(-1, 1).mean(0)
        # self.loss_G_L1 = self.criterionL1(self.fake_B, self.real_B) * self.opt.lambda_A
        # self.loss_G = self.loss_G_GAN + self.loss_G_L1
        self.loss_G = self.loss_G_GAN
        # The Numbers 1
        one = torch.FloatTensor([1]).cuda()
        self.loss_G.backward(one)

    def optimize_parameters(self):
        self.forward()

        self.optimizer_D.zero_grad()
        self.backward_D()
        self.optimizer_D.step()

        self.optimizer_G.zero_grad()
        self.backward_G()
        self.optimizer_G.step()

    def optimize_netD_parameters(self):
        make_trainable(self.netD, True)
        for p in self.netD.parameters():
            p.data.clamp(-0.01, 0.01)
        self.forward()
        self.optimizer_D.zero_grad()
        self.backward_wgan_D()
{'fraction_non_alphanumeric': 0.12108311815199944, 'fraction_numerical': 0.009002721753088142, 'mean_word_length': 2.1494505494505494, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 1, 'word_list_counts': {'cursed_substrings.json': 22, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2191124', 'n_tokens_mistral': 4987, 'n_tokens_neox': 4767, 'n_words': 786}