def __init__(self, dataType='BSD', batch_size=10, args=None):
self.dataType = dataType
if dataType == 'BSD':
self.dataPath = './dataset/'
self.imgList = os.listdir(self.dataPath)
self.batchSize = batch_size
self.len = len(self.imgList)
self.loimgs = torch.zeros((300, 3, 32, 32))
self.midImgs = torch.zeros((300, 3, 64, 64))
self.HDImgs = torch.zeros((300, 3, 128, 128))
self.iter = 0
preprocess = torchTrans.Compose([
torchTrans.ToTensor(),
torchTrans.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])
for i in range(self.len):
imgH = cv2.resize(
mpimg.imread(self.dataPath + self.imgList[i + self.iter]),
(128, 128))[:, :, 0:-1]
imgM = cv2.resize(imgH, (64, 64))
img = cv2.resize(imgM, (32, 32))
imgH = preprocess(imgH)
imgM = preprocess(imgM)
img = preprocess(img)
self.loimgs[i, :, :, :] = img
self.midImgs[i, :, :, :] = imgM
self.HDImgs[i, :, :, :] = imgH
elif dataType == 'CIFAR':
train_gen, dev_gen, test_gen = utils.dataset_iterator(args)
self.batchSize = batch_size
self.gen = utils.inf_train_gen(train_gen)
elif dataType == 'PASCAL':
self.dataPath = './VOCdevkit/VOC2012/'
self.imgList = []
for line in open(self.dataPath + 'ImageSets/Main/trainval.txt'):
self.imgList.append(line[0:-1])
self.batchSize = batch_size
self.len = len(self.imgList)
self.loimgs = torch.zeros((self.len, 3, 32, 32))
self.midImgs = torch.zeros((self.len, 3, 64, 64))
self.HDImgs = torch.zeros((self.len, 3, 128, 128))
self.iter = 0
preprocess = torchTrans.Compose([
torchTrans.ToTensor(),
torchTrans.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])
for i in range(self.len):
imgH = cv2.resize(
mpimg.imread(self.dataPath + 'JPEGImages/' +
self.imgList[i + self.iter] + '.jpg'),
(128, 128))
imgM = cv2.resize(imgH, (64, 64))
img = cv2.resize(imgH, (32, 32))
imgH = preprocess(imgH)
imgM = preprocess(imgM)
img = preprocess(img)
self.loimgs[i, :, :, :] = img
self.midImgs[i, :, :, :] = imgM
self.HDImgs[i, :, :, :] = imgH
def train():
args = load_args()
train_gen, dev_gen, test_gen = utils.dataset_iterator(args)
torch.manual_seed(1)
netG, netD, netE = load_models(args)
# optimizerD = optim.Adam(netD.parameters(), lr=1e-4, betas=(0.5, 0.9))
optimizerG = optim.Adam(netG.parameters(), lr=1e-4, betas=(0.5, 0.9))
vgg_scale = 0.0784 # 1/12.75
mse_criterion = nn.MSELoss()
one = torch.FloatTensor([1]).cuda(0)
mone = (one * -1).cuda(0)
gen = utils.inf_train_gen(train_gen)
""" train SRResNet with MSE """
for iteration in range(1, 20000):
start_time = time.time()
# for p in netD.parameters():
# p.requires_grad = False
for p in netE.parameters():
p.requires_grad = False
netG.zero_grad()
_data = next(gen)
real_data, vgg_data = stack_data(args, _data)
real_data_v = autograd.Variable(real_data)
#Perceptual loss
#vgg_data_v = autograd.Variable(vgg_data)
#vgg_features_real = netE(vgg_data_v)
fake = netG(real_data_v)
#vgg_features_fake = netE(fake)
#diff = vgg_features_fake - vgg_features_real.cuda(0)
#perceptual_loss = vgg_scale * ((diff.pow(2)).sum(3).mean()) # mean(sum(square(diff)))
#perceptual_loss.backward(one)
mse_loss = mse_criterion(fake, real_data_v)
mse_loss.backward(one)
optimizerG.step()
save_dir = './plots/' + args.dataset
plot.plot(save_dir, '/mse cost SRResNet',
np.round(mse_loss.data.cpu().numpy(), 4))
if iteration % 50 == 49:
utils.generate_sr_image(iteration, netG, save_dir, args,
real_data_v)
if (iteration < 5) or (iteration % 50 == 49):
plot.flush()
plot.tick()
if iteration % 5000 == 0:
torch.save(netG.state_dict(), './SRResNet_PL.pt')
for iteration in range(args.epochs):
start_time = time.time()
""" Update AutoEncoder """
for p in netD.parameters():
p.requires_grad = False
netG.zero_grad()
netE.zero_grad()
_data = next(gen)
real_data = stack_data(args, _data)
real_data_v = autograd.Variable(real_data)
encoding = netE(real_data_v)
fake = netG(encoding)
ae_loss = ae_criterion(fake, real_data_v)
ae_loss.backward(one)
optimizerE.step()
optimizerG.step()
""" Update D network """
for p in netD.parameters(): # reset requires_grad
p.requires_grad = True # they are set to False below in netG update
for i in range(5):
_data = next(gen)
real_data = stack_data(args, _data)
real_data_v = autograd.Variable(real_data)
# train with real data
netD.zero_grad()
D_real = netD(real_data_v)
D_real = D_real.mean()
D_real.backward(mone)
# train with fake data
noise = torch.randn(args.batch_size, args.dim).cuda()
noisev = autograd.Variable(noise,
volatile=True) # totally freeze netG
# instead of noise, use image
fake = autograd.Variable(netG(real_data_v).data)
inputv = fake
D_fake = netD(inputv)
D_fake = D_fake.mean()
D_fake.backward(one)
# train with gradient penalty
gradient_penalty = ops.calc_gradient_penalty(
args, netD, real_data_v.data, fake.data)
gradient_penalty.backward()
D_cost = D_fake - D_real + gradient_penalty
Wasserstein_D = D_real - D_fake
optimizerD.step()
# Update generator network (GAN)
# noise = torch.randn(args.batch_size, args.dim).cuda()
# noisev = autograd.Variable(noise)
_data = next(gen)
real_data = stack_data(args, _data)
real_data_v = autograd.Variable(real_data)
# again use real data instead of noise
fake = netG(real_data_v)
G = netD(fake)
G = G.mean()
G.backward(mone)
G_cost = -G
optimizerG.step()
# Write logs and save samples
save_dir = './plots/' + args.dataset
plot.plot(save_dir, '/disc cost', np.round(D_cost.cpu().data.numpy(),
4))
plot.plot(save_dir, '/gen cost', np.round(G_cost.cpu().data.numpy(),
4))
plot.plot(save_dir, '/w1 distance',
np.round(Wasserstein_D.cpu().data.numpy(), 4))
# plot.plot(save_dir, '/ae cost', np.round(ae_loss.data.cpu().numpy(), 4))
# Calculate dev loss and generate samples every 100 iters
if iteration % 100 == 99:
dev_disc_costs = []
for images, _ in dev_gen():
imgs = stack_data(args, images)
imgs_v = autograd.Variable(imgs, volatile=True)
D = netD(imgs_v)
_dev_disc_cost = -D.mean().cpu().data.numpy()
dev_disc_costs.append(_dev_disc_cost)
plot.plot(save_dir, '/dev disc cost',
np.round(np.mean(dev_disc_costs), 4))
# utils.generate_image(iteration, netG, save_dir, args)
# utils.generate_ae_image(iteration, netE, netG, save_dir, args, real_data_v)
utils.generate_sr_image(iteration, netG, save_dir, args,
real_data_v)
# Save logs every 100 iters
if (iteration < 5) or (iteration % 100 == 99):
plot.flush()
plot.tick()
def train():
with torch.cuda.device(1):
args = load_args()
train_gen, dev_gen, test_gen = utils.dataset_iterator(args)
torch.manual_seed(1)
netG = first_layer.FirstG(args).cuda()
SecondG = second_layer.SecondG(args).cuda()
SecondE = second_layer.SecondE(args).cuda()
ThridG = third_layer.ThirdG(args).cuda()
ThridE = third_layer.ThirdE(args).cuda()
ThridD = third_layer.ThirdD(args).cuda()
netG.load_state_dict(torch.load('./1stLayer/1stLayerG71999.model'))
SecondG.load_state_dict(torch.load('./2ndLayer/2ndLayerG71999.model'))
SecondE.load_state_dict(torch.load('./2ndLayer/2ndLayerE71999.model'))
ThridE.load_state_dict(torch.load('./3rdLayer/3rdLayerE10999.model'))
ThridG.load_state_dict(torch.load('./3rdLayer/3rdLayerG10999.model'))
optimizerD = optim.Adam(ThridD.parameters(), lr=1e-4, betas=(0.5, 0.9))
optimizerG = optim.Adam(ThridG.parameters(), lr=1e-4, betas=(0.5, 0.9))
optimizerE = optim.Adam(ThridE.parameters(), lr=1e-4, betas=(0.5, 0.9))
ae_criterion = nn.MSELoss()
one = torch.FloatTensor([1]).cuda()
mone = (one * -1).cuda()
dataLoader = BSDDataLoader(args.dataset, args.batch_size, args)
for iteration in range(args.epochs):
start_time = time.time()
""" Update AutoEncoder """
for p in ThridD.parameters():
p.requires_grad = False
ThridG.zero_grad()
ThridE.zero_grad()
real_data = dataLoader.getNextHDBatch().cuda()
real_data_v = autograd.Variable(real_data)
encoding = ThridE(real_data_v)
fake = ThridG(encoding)
ae_loss = ae_criterion(fake, real_data_v)
ae_loss.backward(one)
optimizerE.step()
optimizerG.step()
""" Update D network """
for p in ThridD.parameters():
p.requires_grad = True
for i in range(5):
real_data = dataLoader.getNextHDBatch().cuda()
real_data_v = autograd.Variable(real_data)
# train with real data
ThridD.zero_grad()
D_real = ThridD(real_data_v)
D_real = D_real.mean()
D_real.backward(mone)
# train with fake data
noise = generateTensor(args.batch_size).cuda()
noisev = autograd.Variable(noise, volatile=True)
fake = autograd.Variable(ThridG(ThridE(SecondG(SecondE(netG(noisev, True), True)), True)).data)
inputv = fake
D_fake = ThridD(inputv)
D_fake = D_fake.mean()
D_fake.backward(one)
# train with gradient penalty
gradient_penalty = ops.calc_gradient_penalty(args,
ThridD, real_data_v.data, fake.data)
gradient_penalty.backward()
optimizerD.step()
# Update generator network (GAN)
noise = generateTensor(args.batch_size).cuda()
noisev = autograd.Variable(noise)
fake = ThridG(ThridE(SecondG(SecondE(netG(noisev, True), True)), True))
G = ThridD(fake)
G = G.mean()
G.backward(mone)
G_cost = -G
optimizerG.step()
# Write logs and save samples
save_dir = './plots/' + args.dataset
# Calculate dev loss and generate samples every 100 iters
if iteration % 1000 == 999:
torch.save(ThridE.state_dict(), './3rdLayer/3rdLayerE%d.model' % iteration)
torch.save(ThridG.state_dict(), './3rdLayer/3rdLayerG%d.model' % iteration)
utils.generate_image(iteration, netG, save_dir, args)
utils.generate_MidImage(iteration, netG, SecondE, SecondG, save_dir, args)
utils.generate_HDImage(iteration, netG, SecondE, SecondG, ThridE, ThridG, save_dir, args)
if iteration % 2000 == 1999:
noise = generateTensor(args.batch_size).cuda()
noisev = autograd.Variable(noise, volatile=True)
fake = autograd.Variable(ThridG(ThridE(SecondG(SecondE(netG(noisev, True), True)), True)).data)
print(inception_score(fake.data.cpu().numpy(), resize=True, batch_size=5)[0])
endtime = time.time()
print('iter:', iteration, 'total time %4f' % (endtime-start_time), 'ae loss %4f' % ae_loss.data[0],
'G cost %4f' % G_cost.data[0])
print('Accuracy at epoch {0}: {1} '.format(epoch, total_correct_preds / total_samples))
print('Took: {0} seconds'.format(time.time() - start_time))
def train(self, train_init, test_init, n_epochs):
writer = tf.summary.FileWriter('graphs/vgg_net', tf.get_default_graph())
with tf.Session() as sess:
sess.run(tf.global_variables_initializer())
self.save_model(sess, 'vgg_face_after_train.data')
step = self.global_step.eval()
for epoch in range(n_epochs):
self.evaluation(sess, test_init, writer, epoch, step)
step = self.train_one_epoch(sess, train_init, writer, epoch, step)
# self.evaluation(sess, test_init, writer, epoch, step)
self.save_model(sess, 'vgg_face_after_train.data')
writer.close()
if __name__ == '__main__':
vgg = VGGFace()
train_set, val_set = utils.load_face_dataset(imgs_path='/srv/node/sdc1/image_data/CelebA/Img/img_align_celeba',
attr_file='/srv/node/sdc1/image_data/CelebA/Anno/list_attr_celeba.txt',
partition_file='/srv/node/sdc1/image_data/CelebA/Eval/list_eval_partition.txt',
cpu_cores=vgg.cpu_cores, batch_size=vgg.batch_size)
train_init, test_init, x, y = utils.dataset_iterator(train_set, val_set)
vgg.load_model(model_path='Weights_imageNet')
vgg.build(x, y)
vgg.train(train_init, test_init, n_epochs=20)
def train():
args = load_args()
train_gen, dev_gen, test_gen = utils.dataset_iterator(args)
torch.manual_seed(1)
np.set_printoptions(precision=4)
netG, netD, netE = load_models(args)
optimizerD = optim.Adam(netD.parameters(), lr=1e-4, betas=(0.5, 0.9))
optimizerG = optim.Adam(netG.parameters(), lr=1e-4, betas=(0.5, 0.9))
optimizerE = optim.Adam(netE.parameters(), lr=1e-4, betas=(0.5, 0.9))
ae_criterion = nn.MSELoss()
one = torch.FloatTensor([1]).cuda()
mone = (one * -1).cuda()
gen = utils.inf_train_gen(train_gen)
preprocess = torchvision.transforms.Compose([
torchvision.transforms.ToTensor(),
torchvision.transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])
for iteration in range(args.epochs):
start_time = time.time()
""" Update AutoEncoder """
for p in netD.parameters():
p.requires_grad = False
netG.zero_grad()
netE.zero_grad()
_data = next(gen)
real_data = stack_data(args, _data)
real_data_v = autograd.Variable(real_data)
encoding = netE(real_data_v)
fake = netG(encoding)
ae_loss = ae_criterion(fake, real_data_v)
ae_loss.backward(one)
optimizerE.step()
optimizerG.step()
""" Update D network """
for p in netD.parameters(): # reset requires_grad
p.requires_grad = True # they are set to False below in netG update
for i in range(5):
_data = next(gen)
real_data = stack_data(args, _data)
real_data_v = autograd.Variable(real_data)
# train with real data
netD.zero_grad()
D_real = netD(real_data_v)
D_real = D_real.mean()
D_real.backward(mone)
# train with fake data
noise = torch.randn(args.batch_size, args.dim).cuda()
noisev = autograd.Variable(noise,
volatile=True) # totally freeze netG
fake = autograd.Variable(netG(noisev).data)
inputv = fake
D_fake = netD(inputv)
D_fake = D_fake.mean()
D_fake.backward(one)
# train with gradient penalty
gradient_penalty = ops.calc_gradient_penalty(
args, netD, real_data_v.data, fake.data)
gradient_penalty.backward()
D_cost = D_fake - D_real + gradient_penalty
Wasserstein_D = D_real - D_fake
optimizerD.step()
# Update generator network (GAN)
noise = torch.randn(args.batch_size, args.dim).cuda()
noisev = autograd.Variable(noise)
fake = netG(noisev)
G = netD(fake)
G = G.mean()
G.backward(mone)
G_cost = -G
optimizerG.step()
# Write logs and save samples
save_dir = './plots/' + args.dataset
plot.plot(save_dir, '/disc cost', np.round(D_cost.cpu().data.numpy(),
4))
plot.plot(save_dir, '/gen cost', np.round(G_cost.cpu().data.numpy(),
4))
plot.plot(save_dir, '/w1 distance',
np.round(Wasserstein_D.cpu().data.numpy(), 4))
plot.plot(save_dir, '/ae cost', np.round(ae_loss.data.cpu().numpy(),
4))
# Calculate dev loss and generate samples every 100 iters
if iteration % 100 == 99:
dev_disc_costs = []
for images, _ in dev_gen():
imgs = stack_data(args, images)
imgs_v = autograd.Variable(imgs, volatile=True)
D = netD(imgs_v)
_dev_disc_cost = -D.mean().cpu().data.numpy()
dev_disc_costs.append(_dev_disc_cost)
plot.plot(save_dir, '/dev disc cost',
np.round(np.mean(dev_disc_costs), 4))
# utils.generate_image(iteration, netG, save_dir, args)
utils.generate_ae_image(iteration, netE, netG, save_dir, args,
real_data_v)
# Save logs every 100 iters
if (iteration < 5) or (iteration % 100 == 99):
plot.flush()
plot.tick()
vgg = VggMultiTask(is_from_vgg_weight=False)
train_set_image_net, test_set_image_net = utils.load_image_net_dataset(
imgs_path='/srv/node/sdc1/image_data/img_val',
label_path='ILSVRC_labels.txt',
cpu_cores=vgg.cpu_cores,
batch_size=vgg.batch_size)
train_set_celeba, test_set_celeba = utils.load_face_dataset(
imgs_path='/srv/node/sdc1/image_data/CelebA/Img/img_align_celeba',
attr_file='/srv/node/sdc1/image_data/CelebA/Anno/list_attr_celeba.txt',
partition_file=
'/srv/node/sdc1/image_data/CelebA/Eval/list_eval_partition.txt',
cpu_cores=vgg.cpu_cores,
batch_size=vgg.batch_size)
train_init_image_net, test_init_image_net, x_image_net, y_image_net = utils.dataset_iterator(
train_set_image_net, test_set_image_net)
train_init_celeba, test_init_celeba, x_celeba, y_celeba = utils.dataset_iterator(
train_set_celeba, test_set_celeba)
vgg.load_model('vgg_result/20181023/vgg_multi_after_train.data')
vgg.build([x_image_net, x_celeba], [y_image_net, y_celeba])
vgg.train(train_init_image_net,
test_init_image_net,
train_init_celeba,
test_init_celeba,
n_epochs=1)
# vgg.load_model('vgg_result/20181023/vgg_multi_after_train.data')
# vgg.build([x_image_net, x_celeba], [y_image_net, y_celeba])
# vgg.evaluation(test_init_image_net, test_init_celeba)
# vgg.train(train_init_image_net, test_init_image_net, train_init_celeba, test_init_celeba, n_epochs=3)
def train():
args = load_args()
train_gen = utils.dataset_iterator(args)
dev_gen = utils.dataset_iterator(args)
torch.manual_seed(1)
netG, netD, netE = load_models(args)
if args.use_spectral_norm:
optimizerD = optim.Adam(filter(lambda p: p.requires_grad,
netD.parameters()), lr=2e-4, betas=(0.0,0.9))
else:
optimizerD = optim.Adam(netD.parameters(), lr=2e-4, betas=(0.5, 0.9))
optimizerG = optim.Adam(netG.parameters(), lr=2e-4, betas=(0.5, 0.9))
optimizerE = optim.Adam(netE.parameters(), lr=2e-4, betas=(0.5, 0.9))
schedulerD = optim.lr_scheduler.ExponentialLR(optimizerD, gamma=0.99)
schedulerG = optim.lr_scheduler.ExponentialLR(optimizerG, gamma=0.99)
schedulerE = optim.lr_scheduler.ExponentialLR(optimizerE, gamma=0.99)
ae_criterion = nn.MSELoss()
one = torch.FloatTensor([1]).cuda()
mone = (one * -1).cuda()
gen = utils.inf_train_gen(train_gen)
for iteration in range(args.epochs):
start_time = time.time()
""" Update AutoEncoder """
for p in netD.parameters():
p.requires_grad = False
netG.zero_grad()
netE.zero_grad()
_data = next(gen)
# real_data = stack_data(args, _data)
real_data = _data
real_data_v = autograd.Variable(real_data).cuda()
encoding = netE(real_data_v)
fake = netG(encoding)
ae_loss = ae_criterion(fake, real_data_v)
ae_loss.backward(one)
optimizerE.step()
optimizerG.step()
""" Update D network """
for p in netD.parameters():
p.requires_grad = True
for i in range(5):
_data = next(gen)
# real_data = stack_data(args, _data)
real_data = _data
real_data_v = autograd.Variable(real_data).cuda()
# train with real data
netD.zero_grad()
D_real = netD(real_data_v)
D_real = D_real.mean()
D_real.backward(mone)
# train with fake data
noise = torch.randn(args.batch_size, args.dim).cuda()
noisev = autograd.Variable(noise, volatile=True)
fake = autograd.Variable(netG(noisev).data)
inputv = fake
D_fake = netD(inputv)
D_fake = D_fake.mean()
D_fake.backward(one)
# train with gradient penalty
gradient_penalty = ops.calc_gradient_penalty(args,
netD, real_data_v.data, fake.data)
gradient_penalty.backward()
D_cost = D_fake - D_real + gradient_penalty
Wasserstein_D = D_real - D_fake
optimizerD.step()
# Update generator network (GAN)
noise = torch.randn(args.batch_size, args.dim).cuda()
noisev = autograd.Variable(noise)
fake = netG(noisev)
G = netD(fake)
G = G.mean()
G.backward(mone)
G_cost = -G
optimizerG.step()
schedulerD.step()
schedulerG.step()
schedulerE.step()
# Write logs and save samples
save_dir = './plots/'+args.dataset
plot.plot(save_dir, '/disc cost', D_cost.cpu().data.numpy())
plot.plot(save_dir, '/gen cost', G_cost.cpu().data.numpy())
plot.plot(save_dir, '/w1 distance', Wasserstein_D.cpu().data.numpy())
plot.plot(save_dir, '/ae cost', ae_loss.data.cpu().numpy())
# Calculate dev loss and generate samples every 100 iters
if iteration % 100 == 99:
dev_disc_costs = []
for i, (images, _) in enumerate(dev_gen):
# imgs = stack_data(args, images)
imgs = images
imgs_v = autograd.Variable(imgs, volatile=True).cuda()
D = netD(imgs_v)
_dev_disc_cost = -D.mean().cpu().data.numpy()
dev_disc_costs.append(_dev_disc_cost)
plot.plot(save_dir ,'/dev disc cost', np.mean(dev_disc_costs))
utils.generate_image(iteration, netG, save_dir, args)
# utils.generate_ae_image(iteration, netE, netG, save_dir, args, real_data_v)
# Save logs every 100 iters
if (iteration < 5) or (iteration % 100 == 99):
plot.flush()
plot.tick()
if iteration % 100 == 0:
utils.save_model(netG, optimizerG, iteration,
'models/{}/G_{}'.format(args.dataset, iteration))
utils.save_model(netD, optimizerD, iteration,
'models/{}/D_{}'.format(args.dataset, iteration))