def evaluate(args):
content_image = utils.tensor_load_rgbimage(args.content_image,
size=args.content_size,
keep_asp=True)
content_image = content_image.unsqueeze(0)
style = utils.tensor_load_rgbimage(args.style_image, size=args.style_size)
style = style.unsqueeze(0)
style = utils.preprocess_batch(style)
vgg = Vgg16()
utils.init_vgg16(args.vgg_model_dir)
vgg.load_state_dict(
torch.load(os.path.join(args.vgg_model_dir, "vgg16.weight")))
style_model = Net()
style_model.load_state_dict(torch.load(args.model))
if args.cuda:
style_model.cuda()
vgg.cuda()
content_image = content_image.cuda()
style = style.cuda()
style_v = Variable(style, volatile=True)
style_v = utils.subtract_imagenet_mean_batch(style_v)
features_style = vgg(style_v)
gram_style = [utils.gram_matrix(y) for y in features_style]
content_image = Variable(utils.preprocess_batch(content_image),
volatile=True)
style_model.setTarget(gram_style[2].data)
output = style_model(content_image)
utils.tensor_save_bgrimage(output.data[0], args.output_image, args.cuda)
def optimize(args):
""" Gatys et al. CVPR 2017
ref: Image Style Transfer Using Convolutional Neural Networks
"""
# load the content and style target
content_image = utils.tensor_load_rgbimage(args.content_image,
size=args.content_size,
keep_asp=True)
content_image = content_image.unsqueeze(0)
content_image = Variable(utils.preprocess_batch(content_image),
requires_grad=False)
content_image = utils.subtract_imagenet_mean_batch(content_image)
style_image = utils.tensor_load_rgbimage(args.style_image,
size=args.style_size)
style_image = style_image.unsqueeze(0)
style_image = Variable(utils.preprocess_batch(style_image),
requires_grad=False)
style_image = utils.subtract_imagenet_mean_batch(style_image)
# load the pre-trained vgg-16 and extract features
vgg = Vgg16()
utils.init_vgg16(args.vgg_model_dir)
vgg.load_state_dict(
torch.load(os.path.join(args.vgg_model_dir, "vgg16.weight")))
if args.cuda:
content_image = content_image.cuda()
style_image = style_image.cuda()
vgg.cuda()
features_content = vgg(content_image)
f_xc_c = Variable(features_content[1].data, requires_grad=False)
features_style = vgg(style_image)
gram_style = [utils.gram_matrix(y) for y in features_style]
# init optimizer
output = Variable(content_image.data, requires_grad=True)
optimizer = Adam([output], lr=args.lr)
mse_loss = torch.nn.MSELoss()
# optimizing the images
for e in range(args.iters):
utils.imagenet_clamp_batch(output, 0, 255)
optimizer.zero_grad()
features_y = vgg(output)
content_loss = args.content_weight * mse_loss(features_y[1], f_xc_c)
style_loss = 0.
for m in range(len(features_y)):
gram_y = utils.gram_matrix(features_y[m])
gram_s = Variable(gram_style[m].data, requires_grad=False)
style_loss += args.style_weight * mse_loss(gram_y, gram_s)
total_loss = content_loss + style_loss
if (e + 1) % args.log_interval == 0:
print(total_loss.data.cpu().numpy()[0])
total_loss.backward()
optimizer.step()
# save the image
output = utils.add_imagenet_mean_batch(output)
utils.tensor_save_bgrimage(output.data[0], args.output_image, args.cuda)
def init_vgg16(model_folder):
from torchvision import models
vgglua = models.vgg16(pretrained=True)
vgglua.eval()
vgg = Vgg16()
for (src, dst) in zip(vgglua.parameters(), vgg.parameters()):
dst.data[:] = src
torch.save(vgg.state_dict(), os.path.join(model_folder, 'vgg16.weight'))
def init_vgg16(model_folder):
"""load the vgg16 model feature"""
if not os.path.exists(os.path.join(model_folder, 'vgg16.weight')):
if not os.path.exists(os.path.join(model_folder, 'vgg16.t7')):
os.system(
'wget http://cs.stanford.edu/people/jcjohns/fast-neural-style/models/vgg16.t7 -O '
+ os.path.join(model_folder, 'vgg16.t7'))
vgglua = torchfile.load(os.path.join(model_folder, 'vgg16.t7'))
vgg = Vgg16()
for (src, dst) in zip(vgglua.parameters()[0], vgg.parameters()):
dst.data[:] = src
torch.save(vgg.state_dict(), os.path.join(model_folder,
'vgg16.weight'))
def extract_feats(args):
vgg = Vgg16()
utils.init_vgg16(args.vgg_model_dir)
cap = dset.CocoCaptions(
root='/Pulsar1/Datasets/coco/train2014/train2014',
annFile=
'/Neutron9/sahil.c/datasets/annotations/captions_train2014.json',
transform=transforms.ToTensor())
print('Number of samples: ', len(cap))
for i, t in cap:
image = i.unsqueeze(0)
image = Variable(utils.preprocess_batch(image), requires_grad=False)
image = utils.subtract_imagenet_mean_batch(image)
features_content = vgg(image)
#haze_train = haze_train.cuda()
#free_train = free_train.cuda()
netD_A.cuda()
#netD_B.cuda()
netG.cuda()
criterionBCE.cuda()
criterionMSE.cuda()
criterionCycle.cuda()
print('done')
lamdaA = opt.lamdaA
lamdaP = opt.lamdaP
# Initialize VGG-16
vgg = Vgg16()
utils.init_vgg16('./models/')
vgg.load_state_dict(torch.load(os.path.join('./models/', "vgg16.weight")))
vgg.cuda()
# pdb.set_trace()
''' set optimizer'''
#optimizerD = optim.Adam(itertools.chain(netD_A.parameters(),netD_B.parameters()), lr = opt.lrD, betas = (opt.beta1, 0.999))
optimizerD = optim.Adam(netD_A.parameters(),
lr=opt.lrD,
betas=(opt.beta1, 0.999))
optimizerG = optim.Adam(netG.parameters(),
lr=opt.lrG,
betas=(opt.beta1, 0.999))
'''test data'''
#val_target = torch.FloatTensor(opt.valBatchSize, outputChannelSize, opt.imageSize, opt.imageSize)
def train(args):
check_paths(args)
np.random.seed(args.seed)
torch.manual_seed(args.seed)
if args.cuda:
torch.cuda.manual_seed(args.seed)
kwargs = {'num_workers': 0, 'pin_memory': False}
else:
kwargs = {}
transform = transforms.Compose([
transforms.Scale(args.image_size),
transforms.CenterCrop(args.image_size),
transforms.ToTensor(),
transforms.Lambda(lambda x: x.mul(255))
])
train_dataset = datasets.ImageFolder(args.dataset, transform)
train_loader = DataLoader(train_dataset,
batch_size=args.batch_size,
**kwargs)
style_model = Net(ngf=args.ngf)
if args.resume is not None:
print('Resuming, initializing using weight from {}.'.format(
args.resume))
style_model.load_state_dict(torch.load(args.resume))
print(style_model)
optimizer = Adam(style_model.parameters(), args.lr)
mse_loss = torch.nn.MSELoss()
vgg = Vgg16()
utils.init_vgg16(args.vgg_model_dir)
vgg.load_state_dict(
torch.load(os.path.join(args.vgg_model_dir, "vgg16.weight")))
if args.cuda:
style_model.cuda()
vgg.cuda()
style_loader = StyleLoader(args.style_folder, args.style_size)
for e in range(args.epochs):
style_model.train()
agg_content_loss = 0.
agg_style_loss = 0.
count = 0
for batch_id, (x, _) in enumerate(train_loader):
n_batch = len(x)
count += n_batch
optimizer.zero_grad()
x = Variable(utils.preprocess_batch(x))
if args.cuda:
x = x.cuda()
style_v = style_loader.get(batch_id)
style_model.setTarget(style_v)
style_v = utils.subtract_imagenet_mean_batch(style_v)
features_style = vgg(style_v)
gram_style = [utils.gram_matrix(y) for y in features_style]
y = style_model(x)
xc = Variable(x.data.clone(), volatile=True)
y = utils.subtract_imagenet_mean_batch(y)
xc = utils.subtract_imagenet_mean_batch(xc)
features_y = vgg(y)
features_xc = vgg(xc)
f_xc_c = Variable(features_xc[1].data, requires_grad=False)
content_loss = args.content_weight * mse_loss(
features_y[1], f_xc_c)
style_loss = 0.
for m in range(len(features_y)):
gram_y = utils.gram_matrix(features_y[m])
gram_s = Variable(gram_style[m].data,
requires_grad=False).repeat(
args.batch_size, 1, 1, 1)
style_loss += args.style_weight * mse_loss(
gram_y, gram_s[:n_batch, :, :])
total_loss = content_loss + style_loss
total_loss.backward()
optimizer.step()
agg_content_loss += content_loss.data[0]
agg_style_loss += style_loss.data[0]
if (batch_id + 1) % args.log_interval == 0:
mesg = "{}\tEpoch {}:\t[{}/{}]\tcontent: {:.6f}\tstyle: {:.6f}\ttotal: {:.6f}".format(
time.ctime(), e + 1, count, len(train_dataset),
agg_content_loss / (batch_id + 1),
agg_style_loss / (batch_id + 1),
(agg_content_loss + agg_style_loss) / (batch_id + 1))
print(mesg)
if (batch_id + 1) % (4 * args.log_interval) == 0:
# save model
style_model.eval()
style_model.cpu()
save_model_filename = "Epoch_" + str(e) + "iters_" + str(
count) + "_" + str(time.ctime()).replace(
' ', '_') + "_" + str(args.content_weight) + "_" + str(
args.style_weight) + ".model"
save_model_path = os.path.join(args.save_model_dir,
save_model_filename)
torch.save(style_model.state_dict(), save_model_path)
style_model.train()
style_model.cuda()
print("\nCheckpoint, trained model saved at", save_model_path)
# save model
style_model.eval()
style_model.cpu()
save_model_filename = "Final_epoch_" + str(args.epochs) + "_" + str(
time.ctime()).replace(' ', '_') + "_" + str(
args.content_weight) + "_" + str(args.style_weight) + ".model"
save_model_path = os.path.join(args.save_model_dir, save_model_filename)
torch.save(style_model.state_dict(), save_model_path)
print("\nDone, trained model saved at", save_model_path)
def train(print_every=10):
checkpaths(opt)
train_set = DatasetFromFolder(opt, True)
test_set = DatasetFromFolder(opt, False)
training_data_loader = DataLoader(dataset=train_set, num_workers=opt.threads, batch_size=opt.batchSize, shuffle=True)
testing_data_loader = DataLoader(dataset=test_set, num_workers=opt.threads, batch_size=1, shuffle=False)
norm_layer = get_norm_layer(norm_type='batch')
netD = NLayerDiscriminator(opt.input_nc, opt.ndf, n_layers=1, norm_layer=norm_layer,use_sigmoid=False, gpu_ids=opt.gpu_ids)
netG = MyUnetGenerator(opt.input_nc, opt.output_nc, 8, opt.ngf, norm_layer=norm_layer, use_dropout=False, gpu_ids=opt.gpu_ids)
netE = MyEncoder(opt.input_nc, opt.output_nc, 8, opt.ngf, norm_layer=norm_layer,use_dropout=False, gpu_ids=opt.gpu_ids)
netVGG = Vgg16()
# utils.init_vgg16(opt.model_dir)
netVGG.load_state_dict(torch.load(os.path.join(opt.model_dir, "vgg16.weight")))
VGG = make_encoder(model_file=opt.model_vgg)
perceptual_loss = PerceptualLoss(VGG, 3)
VGG.cuda()
netG.cuda()
netD.cuda()
netE.cuda()
netG.apply(weights_init)
netD.apply(weights_init)
netE.apply(weights_init)
criterionGAN = GANLoss(use_lsgan=not opt.no_lsgan)
criterionL1 = torch.nn.L1Loss()
mse_loss = torch.nn.MSELoss()
criterionCEL = nn.CrossEntropyLoss()
# initialize optimizers
optimizer_G = torch.optim.Adam(netG.parameters(), lr=opt.lr, betas=(opt.beta1, 0.999))
optimizer_D = torch.optim.Adam(netD.parameters(), lr=opt.lr, betas=(opt.beta1, 0.999))
optimizer_E = torch.optim.Adam(netE.parameters(), lr=opt.lr, betas=(opt.beta1, 0.999))
print('=========== Networks initialized ============')
print_network(netG)
print_network(netD)
print('=============================================')
f = open('./checkpoint/loss.txt', 'w')
f2 = open('./checkpoint/recognition.txt', 'w')
strat_time = time.time()
for epoch in range(1, opt.n_epoch + 1):
D_running_loss = 0.0
G_running_loss = 0.0
G2_running_loss = 0.0
for (i, batch) in enumerate(training_data_loader, 1):
real_p, real_s, identity = Variable(batch[0]), Variable(batch[1]), Variable(batch[2].squeeze(1))
location = batch[3]
real_p, real_s, identity = real_p.cuda(), real_s.cuda(), identity.cuda()
optimizer_D.zero_grad()
# fake
parsing_feature = netE.forward(real_p[:, 3:, :, :])
fake_s = netG.forward(real_p[:, 0:3, :, :], parsing_feature)
fake_ps = torch.cat((fake_s, real_p), 1)
pred_fake = netD.forward(fake_ps.detach())
loss_D_fake = criterionGAN(pred_fake, False)
# real
real_ps = torch.cat((real_s, real_p), 1)
pred_real = netD.forward(real_ps)
loss_D_real = criterionGAN(pred_real, True)
loss_D = (loss_D_real + loss_D_fake) * 0.5
loss_D.backward()
optimizer_D.step()
optimizer_G.zero_grad()
optimizer_E.zero_grad()
pred_fake = netD.forward(fake_ps)
loss_G_GAN = criterionGAN(pred_fake, True)
# loss_G_L1 = criterionL1(fake_s, real_s) * opt.lambda1
# !!!!!!!-------- a2b need modified cirterionL1 -----------------!!!
loss_global = criterionL1(fake_s, real_s)
loss_local = localLossL1(fake_s, real_s, real_p, criterionL1)
loss_G_L1 = opt.alpha1 * loss_global + (1 - opt.alpha1) * loss_local
loss_G_L1 *= opt.lambda1
b,c,w,h = fake_s.shape
yh = fake_s.expand(b,3,w,h)
ys = real_s.expand(b,3,w,h)
_mean = Variable(torch.Tensor([0.485, 0.456, 0.406]).view(1, 3, 1, 1).expand_as(yh)).cuda()
_var = Variable(torch.Tensor([0.229, 0.224, 0.225]).view(1, 3, 1, 1).expand_as(yh)).cuda()
yh = yh / 2 + 0.5
ys = ys / 2 + 0.5
yh = (yh - _mean) / _var
ys = (ys - _mean) / _var
loss_recog = perceptual_loss(yh, ys)
loss_G = loss_G_GAN + loss_G_L1 + opt.styleParam * loss_recog
loss_G.backward()
optimizer_G.step()
optimizer_E.step()
'''======================================================================='''
D_running_loss += loss_D.data[0]
G_running_loss += loss_G.data[0]
G2_running_loss += loss_G.data[0]
if i % print_every == 0:
end_time = time.time()
time_delta = usedtime(strat_time, end_time)
print('[%s-%d, %5d] D loss: %.3f ; G loss: %.3f' % (time_delta, epoch, i + 1, D_running_loss / print_every, G_running_loss / print_every))
f.write('%d,%d,D_loss:%.5f,GAN_loss:%.5f,L1Loss:%.5f\r\n' % (epoch, i + 1, loss_D.data[0], loss_G_GAN.data[0],loss_G_L1.data[0]))
f2.write('%d,%d,loss_recog_loss:%.5f\r\n' % (epoch, i + 1, loss_recog.data[0]))
D_running_loss = 0.0
G_running_loss = 0.0
G2_running_loss = 0.0
f.flush()
f2.flush()
if epoch >= 500 and epoch % 50 == 0:
test(epoch, netG, netE, testing_data_loader, opt)
checkpoint(epoch, netD, netG, netE)
f.close()
f2.close()