<commit_before>import argparse
import os
import random
from math import log10
import scipy.stats as stats
import torch.backends.cudnn as cudnn
import torch.optim as optim
import torch.optim.lr_scheduler as lr_scheduler
import torchvision.utils as vutils
from tensorboardX import SummaryWriter
from torch.autograd import Variable, grad
from models.cp_model import *
from data.proData import CreateDataLoader

parser = argparse.ArgumentParser()
parser.add_argument('--datarootC', required=True, help='path to colored dataset')
parser.add_argument('--datarootS', required=True, help='path to sketch dataset')
parser.add_argument('--workers', type=int, help='number of data loading workers', default=4)
parser.add_argument('--batchSize', type=int, default=16, help='input batch size')
parser.add_argument('--imageSize', type=int, default=512, help='the height / width of the input image to network')
parser.add_argument('--cut', type=int, default=1, help='cut backup frequency')
parser.add_argument('--niter', type=int, default=700, help='number of epochs to train for')
parser.add_argument('--ngf', type=int, default=64)
parser.add_argument('--ndf', type=int, default=64)
parser.add_argument('--lrG', type=float, default=0.0001, help='learning rate, default=0.0001')
parser.add_argument('--lrD', type=float, default=0.0001, help='learning rate, default=0.0001')
parser.add_argument('--beta1', type=float, default=0.5, help='beta1 for adam. default=0.5')
parser.add_argument('--cuda', action='store_true', help='enables cuda')
parser.add_argument('--netG', default='', help="path to netG (to continue training)")
parser.add_argument('--netD', default='', help="path to netD (to continue training)")
parser.add_argument('--netD2', default='', help="path to netD2 (to continue training)")
parser.add_argument('--optim', action='store_true', help='load optimizer\'s checkpoint')
parser.add_argument('--outf', default='.', help='folder to output images and model checkpoints')
parser.add_argument('--optf', default='.', help='folder to optimizer checkpoints')
parser.add_argument('--Diters', type=int, default=1, help='number of D iters per each G iter')
parser.add_argument('--manualSeed', type=int, default=2345, help='random seed to use. Default=1234')
parser.add_argument('--baseGeni', type=int, default=2500, help='start base of pure pair L1 loss')
parser.add_argument('--geni', type=int, default=0, help='continue gen image num')
parser.add_argument('--epoi', type=int, default=0, help='continue epoch num')
parser.add_argument('--env', type=str, default=None, help='tensorboard env')
parser.add_argument('--advW', type=float, default=0.01, help='adversarial weight, default=0.01')
parser.add_argument('--gpW', type=float, default=10, help='gradient penalty weight')
parser.add_argument('--drift', type=float, default=0.001, help='wasserstein drift weight')
parser.add_argument('--mseW', type=float, default=0.01, help='MSE loss weight')
parser.add_argument('--MSE', action='store_true', help='enables pure MSE')
parser.add_argument('--feat', action='store_true', help='enables feat test')
parser.add_argument('--cp', action='store_true', help='enables comp train')
parser.add_argument('--cp2', action='store_true', help='enables comp train')
parser.add_argument('--cp3', action='store_true', help='enables comp train')
parser.add_argument('--cp4', action='store_true', help='enables comp train')

opt = parser.parse_args()
print(opt)

####### regular set up
if torch.cuda.is_available() and not opt.cuda:
    print("WARNING: You have a CUDA device, so you should probably run with --cuda")
gen_iterations = opt.geni
try:
    os.makedirs(opt.outf)
except OSError:
    pass
# random seed setup                                  # !!!!!
print("Random Seed: ", opt.manualSeed)
random.seed(opt.manualSeed)
torch.manual_seed(opt.manualSeed)
if opt.cuda:
    torch.cuda.manual_seed(opt.manualSeed)
cudnn.benchmark = True
####### regular set up end

writer = SummaryWriter(log_dir=opt.env, comment='this is great')

dataloader = CreateDataLoader(opt)

netG = def_netG(ngf=opt.ngf)
if opt.netG != '':
    netG.load_state_dict(torch.load(opt.netG))
print(netG)

netD = def_netD(ndf=opt.ndf)
netD2 = def_netD(ndf=opt.ndf)
if opt.netD != '':
    netD.load_state_dict(torch.load(opt.netD))
if opt.netD2 != '':
    netD2.load_state_dict(torch.load(opt.netD2))
print(netD)

netF = def_netF()
print(netF)

criterion_L1 = nn.L1Loss()
criterion_MSE = nn.MSELoss()
L2_dist = nn.PairwiseDistance(2)
one = torch.FloatTensor([1])
mone = one * -1

fixed_sketch = torch.FloatTensor()
fixed_hint = torch.FloatTensor()
saber = torch.FloatTensor([0.485 - 0.5, 0.456 - 0.5, 0.406 - 0.5]).view(1, 3, 1, 1)
diver = torch.FloatTensor([0.229, 0.224, 0.225]).view(1, 3, 1, 1)

if opt.cuda:
    netD = torch.nn.DataParallel(netD).cuda()
    netG = torch.nn.DataParallel(netG).cuda()
    netF = torch.nn.DataParallel(netF).cuda()
    fixed_sketch, fixed_hint = fixed_sketch.cuda(), fixed_hint.cuda()
    saber, diver = saber.cuda(), diver.cuda()
    criterion_L1 = criterion_L1.cuda()
    criterion_MSE = criterion_MSE.cuda()
    one, mone = one.cuda(), mone.cuda()

if opt.feat:
    netF2 = torch.nn.DataParallel(def_netF2()).cuda()

# setup optimizer
optimizerG = optim.Adam(netG.parameters(), lr=opt.lrG, betas=(opt.beta1, 0.9))
optimizerD = optim.Adam(netD.parameters(), lr=opt.lrD, betas=(opt.beta1, 0.9))
optimizerD2 = optim.Adam(netD2.parameters(), lr=opt.lrD, betas=(opt.beta1, 0.9))

if opt.optim:
    optimizerG.load_state_dict(torch.load('%s/optimG_checkpoint.pth' % opt.optf))
    optimizerD.load_state_dict(torch.load('%s/optimD_checkpoint.pth' % opt.optf))
    optimizerD2.load_state_dict(torch.load('%s/optimD2_checkpoint.pth' % opt.optf))


# schedulerG = lr_scheduler.ReduceLROnPlateau(optimizerG, mode='max', verbose=True, min_lr=0.0000005,
#                                             patience=8)  # 1.5*10^5 iter
# schedulerD = lr_scheduler.ReduceLROnPlateau(optimizerD, mode='max', verbose=True, min_lr=0.0000005,
#                                             patience=8)  # 1.5*10^5 iter

# schedulerG = lr_scheduler.StepLR(optimizerG, step_size=150000, gamma=0.1, last_epoch=-1)  # 1.5*10^5 iter
# schedulerD = lr_scheduler.StepLR(optimizerG, step_size=150000, gamma=0.1, last_epoch=-1)  # 1.5*10^5 iter


# schedulerG = lr_scheduler.MultiStepLR(optimizerG, milestones=[60, 120], gamma=0.1)  # 1.5*10^5 iter
# schedulerD = lr_scheduler.MultiStepLR(optimizerD, milestones=[60, 120], gamma=0.1)


def calc_gradient_penalty(netD, real_data, fake_data, sketch):
    alpha = torch.rand(opt.batchSize, 1, 1, 1)
    alpha = alpha.cuda() if opt.cuda else alpha

    interpolates = alpha * real_data + ((1 - alpha) * fake_data)

    if opt.cuda:
        interpolates = interpolates.cuda()
    interpolates = Variable(interpolates, requires_grad=True)

    disc_interpolates = netD(interpolates, Variable(sketch))[0]
<commit_msg>:bug: fix size mismatch bug<commit_after>import argparse
import os
import random
from math import log10
import scipy.stats as stats
import torch.backends.cudnn as cudnn
import torch.optim as optim
import torch.optim.lr_scheduler as lr_scheduler
import torchvision.utils as vutils
from tensorboardX import SummaryWriter
from torch.autograd import Variable, grad
from models.cp_model import *
from data.proData import CreateDataLoader

parser = argparse.ArgumentParser()
parser.add_argument('--datarootC', required=True, help='path to colored dataset')
parser.add_argument('--datarootS', required=True, help='path to sketch dataset')
parser.add_argument('--workers', type=int, help='number of data loading workers', default=4)
parser.add_argument('--batchSize', type=int, default=16, help='input batch size')
parser.add_argument('--imageSize', type=int, default=512, help='the height / width of the input image to network')
parser.add_argument('--cut', type=int, default=1, help='cut backup frequency')
parser.add_argument('--niter', type=int, default=700, help='number of epochs to train for')
parser.add_argument('--ngf', type=int, default=64)
parser.add_argument('--ndf', type=int, default=64)
parser.add_argument('--lrG', type=float, default=0.0001, help='learning rate, default=0.0001')
parser.add_argument('--lrD', type=float, default=0.0001, help='learning rate, default=0.0001')
parser.add_argument('--beta1', type=float, default=0.5, help='beta1 for adam. default=0.5')
parser.add_argument('--cuda', action='store_true', help='enables cuda')
parser.add_argument('--netG', default='', help="path to netG (to continue training)")
parser.add_argument('--netD', default='', help="path to netD (to continue training)")
parser.add_argument('--netD2', default='', help="path to netD2 (to continue training)")
parser.add_argument('--optim', action='store_true', help='load optimizer\'s checkpoint')
parser.add_argument('--outf', default='.', help='folder to output images and model checkpoints')
parser.add_argument('--optf', default='.', help='folder to optimizer checkpoints')
parser.add_argument('--Diters', type=int, default=1, help='number of D iters per each G iter')
parser.add_argument('--manualSeed', type=int, default=2345, help='random seed to use. Default=1234')
parser.add_argument('--baseGeni', type=int, default=2500, help='start base of pure pair L1 loss')
parser.add_argument('--geni', type=int, default=0, help='continue gen image num')
parser.add_argument('--epoi', type=int, default=0, help='continue epoch num')
parser.add_argument('--env', type=str, default=None, help='tensorboard env')
parser.add_argument('--advW', type=float, default=0.01, help='adversarial weight, default=0.01')
parser.add_argument('--gpW', type=float, default=10, help='gradient penalty weight')
parser.add_argument('--drift', type=float, default=0.001, help='wasserstein drift weight')
parser.add_argument('--mseW', type=float, default=0.01, help='MSE loss weight')
parser.add_argument('--MSE', action='store_true', help='enables pure MSE')
parser.add_argument('--feat', action='store_true', help='enables feat test')
parser.add_argument('--cp', action='store_true', help='enables comp train')
parser.add_argument('--cp2', action='store_true', help='enables comp train')
parser.add_argument('--cp3', action='store_true', help='enables comp train')
parser.add_argument('--cp4', action='store_true', help='enables comp train')

opt = parser.parse_args()
print(opt)

####### regular set up
if torch.cuda.is_available() and not opt.cuda:
    print("WARNING: You have a CUDA device, so you should probably run with --cuda")
gen_iterations = opt.geni
try:
    os.makedirs(opt.outf)
except OSError:
    pass
# random seed setup                                  # !!!!!
print("Random Seed: ", opt.manualSeed)
random.seed(opt.manualSeed)
torch.manual_seed(opt.manualSeed)
if opt.cuda:
    torch.cuda.manual_seed(opt.manualSeed)
cudnn.benchmark = True
####### regular set up end

writer = SummaryWriter(log_dir=opt.env, comment='this is great')

dataloader = CreateDataLoader(opt)

netG = def_netG(ngf=opt.ngf)
if opt.netG != '':
    netG.load_state_dict(torch.load(opt.netG))
print(netG)

netD = def_netD(ndf=opt.ndf)
netD2 = def_netD(ndf=opt.ndf)
if opt.netD != '':
    netD.load_state_dict(torch.load(opt.netD))
if opt.netD2 != '':
    netD2.load_state_dict(torch.load(opt.netD2))
print(netD)

netF = def_netF()
print(netF)

criterion_L1 = nn.L1Loss()
criterion_MSE = nn.MSELoss()
L2_dist = nn.PairwiseDistance(2)
one = torch.FloatTensor([1])
mone = one * -1

fixed_sketch = torch.FloatTensor()
fixed_hint = torch.FloatTensor()
saber = torch.FloatTensor([0.485 - 0.5, 0.456 - 0.5, 0.406 - 0.5]).view(1, 3, 1, 1)
diver = torch.FloatTensor([0.229, 0.224, 0.225]).view(1, 3, 1, 1)

if opt.cuda:
    netD = torch.nn.DataParallel(netD).cuda()
    netG = torch.nn.DataParallel(netG).cuda()
    netF = torch.nn.DataParallel(netF).cuda()
    fixed_sketch, fixed_hint = fixed_sketch.cuda(), fixed_hint.cuda()
    saber, diver = saber.cuda(), diver.cuda()
    criterion_L1 = criterion_L1.cuda()
    criterion_MSE = criterion_MSE.cuda()
    one, mone = one.cuda(), mone.cuda()

if opt.feat:
    netF2 = torch.nn.DataParallel(def_netF2()).cuda()

# setup optimizer
optimizerG = optim.Adam(netG.parameters(), lr=opt.lrG, betas=(opt.beta1, 0.9))
optimizerD = optim.Adam(netD.parameters(), lr=opt.lrD, betas=(opt.beta1, 0.9))
optimizerD2 = optim.Adam(netD2.parameters(), lr=opt.lrD, betas=(opt.beta1, 0.9))

if opt.optim:
    optimizerG.load_state_dict(torch.load('%s/optimG_checkpoint.pth' % opt.optf))
    optimizerD.load_state_dict(torch.load('%s/optimD_checkpoint.pth' % opt.optf))
    optimizerD2.load_state_dict(torch.load('%s/optimD2_checkpoint.pth' % opt.optf))


# schedulerG = lr_scheduler.ReduceLROnPlateau(optimizerG, mode='max', verbose=True, min_lr=0.0000005,
#                                             patience=8)  # 1.5*10^5 iter
# schedulerD = lr_scheduler.ReduceLROnPlateau(optimizerD, mode='max', verbose=True, min_lr=0.0000005,
#                                             patience=8)  # 1.5*10^5 iter

# schedulerG = lr_scheduler.StepLR(optimizerG, step_size=150000, gamma=0.1, last_epoch=-1)  # 1.5*10^5 iter
# schedulerD = lr_scheduler.StepLR(optimizerG, step_size=150000, gamma=0.1, last_epoch=-1)  # 1.5*10^5 iter


# schedulerG = lr_scheduler.MultiStepLR(optimizerG, milestones=[60, 120], gamma=0.1)  # 1.5*10^5 iter
# schedulerD = lr_scheduler.MultiStepLR(optimizerD, milestones=[60, 120], gamma=0.1)


def calc_gradient_penalty(netD, real_data, fake_data, sketch):
    alpha = torch.rand(opt.batchSize // 2, 1, 1, 1)
    alpha = alpha.cuda() if opt.cuda else alpha

    interpolates = alpha * real_data + ((1 - alpha) * fake_data)

    if opt.cuda:
        interpolates = interpolates.cuda()
    interpolates = Variable(interpolates, requires_grad=True)

    disc_interpolates = netD(interpolates, Variable(sketch))[0]

{'fraction_non_alphanumeric': 0.1592075771816933, 'fraction_numerical': 0.03231870435977153, 'mean_word_length': 5.161247216035635, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 10, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5153222', 'n_tokens_mistral': 5186, 'n_tokens_neox': 4903, 'n_words': 1023}