def train(model, img, art, photo, epoch_num, device, content_name_list,
style_name_list):
args = arg_parser()
features = vgg19_features(model, content_name_list, style_name_list,
device)
optimizer = torch.optim.SGD([img.requires_grad_()],
lr=args.lr,
momentum=args.momentum)
_, art_style = features.extract_features(art)
art_style = [i_style.detach() for i_style in art_style]
photo_content, _ = features.extract_features(photo)
photo_content = [i_content.detach() for i_content in photo_content]
for epoch in range(epoch_num):
end_time = time.time()
img_content, img_style = features.extract_features(img)
C_loss = content_loss(img_content, photo_content)
S_loss = style_loss(img_style, art_style)
loss = C_loss * args.content_weight + S_loss
optimizer.zero_grad()
loss.backward()
optimizer.step()
if epoch % args.log == 0:
print('[{0}/{1}]\ttime:{time:.2f}\tloss:{loss:.4f}'.format(epoch, epoch_num,\
time=time.time()-end_time, loss=loss.item()*1e6))
print(C_loss.item(), S_loss.item())
if epoch % args.save_fre == 0:
save_img(epoch, img)
img.data.clamp_(0, 1)
return img
vgg_net = vgg.Model(model_path, width, height)
# get style layer from constant network
network = vgg_net.build(style_image, 0)
style_layer = [
sess.run(network['conv' + str(i) + '_1']) for i in range(1, 6)
]
# get content layer from constant network
network = vgg_net.build(content_image, 0)
content_layer = sess.run(network['conv4_2'])
# style transfer network
network = vgg_net.build(pred_image, 1)
pred_style = [network['conv' + str(i) + '_1'] for i in range(1, 6)]
pred_content = network['conv4_2']
style_loss = loss.style_loss(style_layer, pred_style)
content_loss = loss.content_loss(content_layer, pred_content)
total_loss = args.ALPHA * content_loss + args.BETA * style_loss
default_vars = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES)
vgg_vars = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES,
scope='vggnet')
optimizer = tf.train.AdamOptimizer(args.learning_rate).minimize(
loss=total_loss, var_list=default_vars + vgg_vars)
saver = tf.train.Saver()
# train
print('Training Start !!!')
def optimize():
MODEL_DIR_NAME = os.path.dirname(FLAGS.MODEL_PATH)
if not os.path.exists(MODEL_DIR_NAME):
os.mkdir(MODEL_DIR_NAME)
style_paths = FLAGS.STYLE_IMAGES.split(',')
style_layers = FLAGS.STYLE_LAYERS.split(',')
content_layers = FLAGS.CONTENT_LAYERS.split(',')
# style gram matrix
style_features_t = loss.get_style_features(style_paths, style_layers,
FLAGS.IMAGE_SIZE, FLAGS.STYLE_SCALE, FLAGS.VGG_PATH)
with tf.Graph().as_default(), tf.Session() as sess:
# train_images
images = reader.image(FLAGS.BATCH_SIZE, FLAGS.IMAGE_SIZE,
FLAGS.TRAIN_IMAGES_FOLDER, FLAGS.EPOCHS)
generated = transform.net(images - vgg.MEAN_PIXEL, training=True)
net, _ = vgg.net(FLAGS.VGG_PATH, tf.concat([generated, images], 0) - vgg.MEAN_PIXEL)
# 损失函数
content_loss = loss.content_loss(net, content_layers)
style_loss = loss.style_loss(
net, style_features_t, style_layers) / len(style_paths)
total_loss = FLAGS.STYLE_WEIGHT * style_loss + FLAGS.CONTENT_WEIGHT * content_loss + \
FLAGS.TV_WEIGHT * loss.total_variation_loss(generated)
# 准备训练
global_step = tf.Variable(0, name="global_step", trainable=False)
variable_to_train = []
for variable in tf.trainable_variables():
if not variable.name.startswith('vgg19'):
variable_to_train.append(variable)
train_op = tf.train.AdamOptimizer(FLAGS.LEARNING_RATE).minimize(
total_loss, global_step=global_step, var_list=variable_to_train)
variables_to_restore = []
for v in tf.global_variables():
if not v.name.startswith('vgg19'):
variables_to_restore.append(v)
# 开始训练
saver = tf.train.Saver(variables_to_restore,
write_version=tf.train.SaverDef.V1)
sess.run([tf.global_variables_initializer(),
tf.local_variables_initializer()])
# 加载检查点
ckpt = tf.train.latest_checkpoint(MODEL_DIR_NAME)
if ckpt:
tf.logging.info('Restoring model from {}'.format(ckpt))
saver.restore(sess, ckpt)
coord = tf.train.Coordinator()
threads = tf.train.start_queue_runners(coord=coord)
start_time = time.time()
try:
while not coord.should_stop():
_, loss_t, step = sess.run([train_op, total_loss, global_step])
elapsed_time = time.time() - start_time
start_time = time.time()
if step % 10 == 0:
tf.logging.info(
'step: %d, total loss %f, secs/step: %f' % (step, loss_t, elapsed_time))
if step % 10000 == 0:
saver.save(sess, FLAGS.MODEL_PATH, global_step=step)
tf.logging.info('Save model')
except tf.errors.OutOfRangeError:
saver.save(sess, FLAGS.MODEL_PATH + '-done')
tf.logging.info('Done training -- epoch limit reached')
finally:
coord.request_stop()
coord.join(threads)
def style_transfer(content_img_path,
img_size,
style_img_path,
style_size,
content_layer,
content_weight,
style_layers,
style_weights,
tv_weight,
init_random=False):
"""Perform style transfer from style image to source content image
Args:
content_img_path (str): File location of the content image.
img_size (int): Size of the smallest content image dimension.
style_img_path (str): File location of the style image.
style_size (int): Size of the smallest style image dimension.
content_layer (int): Index of the layer to use for content loss.
content_weight (float): Scalar weight for content loss.
style_layers ([]int): Indices of layers to use for style loss.
style_weights ([]float): List of scalar weights to use for each layer in style_layers.
tv_weigh (float): Scalar weight of total variation regularization term.
init_random (boolean): Whether to initialize the starting image to uniform random noise.
"""
tf.reset_default_graph()
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
sess = tf.Session(config=config)
try:
model = SqueezeNet(ckpt_path=CKPT_PATH, sess=sess)
except NotFoundError:
raise ValueError('checkpoint file is not found, please check %s' %
CKPT_PATH)
# Extract features from content image
content_img = preprocess_image(load_image(content_img_path, size=img_size))
content_feats = model.extract_features(model.image)
# Create content target
content_target = sess.run(content_feats[content_layer],
{model.image: content_img[None]})
# Extract features from style image
style_img = preprocess_image(load_image(style_img_path, size=style_size))
style_feats_by_layer = [content_feats[i] for i in style_layers]
# Create style targets
style_targets = []
for style_feats in style_feats_by_layer:
style_targets.append(gram_matrix(style_feats))
style_targets = sess.run(style_targets, {model.image: style_img[None]})
if init_random:
generated_img = tf.Variable(tf.random_uniform(content_img[None].shape,
0, 1),
name="image")
else:
generated_img = tf.Variable(content_img[None], name="image")
# Extract features from generated image
current_feats = model.extract_features(generated_img)
loss = content_loss(content_weight, current_feats[content_layer], content_target) + \
style_loss(current_feats, style_layers, style_targets, style_weights) + \
total_variation_loss(generated_img, tv_weight)
# Set up optimization parameters
init_learning_rate = 3.0
decayed_learning_rate = 0.1
max_iter = 200
learning_rate = tf.Variable(init_learning_rate, name="lr")
with tf.variable_scope("optimizer") as opt_scope:
train_op = tf.train.AdamOptimizer(learning_rate).minimize(
loss, var_list=[generated_img])
opt_vars = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES,
scope=opt_scope.name)
sess.run(
tf.variables_initializer([learning_rate, generated_img] + opt_vars))
# Create an op that will clamp the image values when run
clamp_image_op = tf.assign(generated_img,
tf.clip_by_value(generated_img, -1.5, 1.5))
display_content_and_style(content_img, style_img)
for t in range(max_iter):
sess.run(train_op)
if t < int(0.90 * max_iter):
sess.run(clamp_image_op)
elif t == int(0.90 * max_iter):
sess.run(tf.assign(learning_rate, decayed_learning_rate))
if t % 20 == 0:
current_loss = sess.run(loss)
print 'Iteration %d: %f' % (t, current_loss)
img = sess.run(generated_img)
plt.imshow(deprocess_image(img[0], rescale=True))
plt.axis('off')
plt.show()
def train(options):
content_weight = options.content_weight
tv_weight = options.tv_weight
initial_lr = options.initial_lr
max_iter = options.max_iter
style_weights = options.style_weights
print_iterations = options.print_iterations
img_size = options.img_size
content = options.content
style = options.style
output = options.output
beta1 = options.beta1
beta2 = options.beta2
epsilon = options.epsilon
h5_file = options.h5_file
content_layer = 12
style_layers = [0, 3, 6, 11, 16]
style_target_vars = []
print(tf.test.is_gpu_available(cuda_only=True))
contentImg = pre_img.load_image(content, size=img_size)
contentImg = pre_img.preprocess_image(contentImg)
styleImg = pre_img.load_image(style, size=img_size)
styleImg = pre_img.preprocess_image(styleImg)
img_var = tf.Variable(contentImg[None], name="image", dtype=tf.float32)
lr_var = tf.Variable(initial_lr, name="lr")
new_img_feats = vgg.extract_features(img_var, h5_file)
content_img_feats = vgg.extract_features(contentImg[None], h5_file)
style_img_feats = vgg.extract_features(styleImg[None], h5_file)
for idx in style_layers:
style_target_vars.append(loss.gram_matrix(style_img_feats[idx]))
with tf.Session(config=tf.ConfigProto(log_device_placement=True)) as sess:
style_loss = loss.style_loss(new_img_feats, style_layers,
style_target_vars, style_weights)
content_loss = loss.content_loss(content_weight,
new_img_feats[content_layer],
content_img_feats[content_layer])
tv_loss = loss.tv_loss(img_var, tv_weight)
total_loss = style_loss + content_loss + tv_loss
optimizer = tf.train.AdamOptimizer(learning_rate=lr_var,
beta1=beta1,
beta2=beta2,
epsilon=epsilon)
training_op = optimizer.minimize(total_loss, var_list=[img_var])
init = tf.global_variables_initializer()
sess.run(init)
sess.run(img_var.initializer)
for t in range(max_iter):
if print_iterations is not None and t % print_iterations == 0:
new_image = img_var.eval()
imageio.imwrite(output + '\\iteration_' + str(t) + '.jpg',
pre_img.deprocess_image(new_image[0]))
sess.run(training_op)
loss_val = sess.run(total_loss)
s_loss = sess.run(style_loss)
c_loss = sess.run(content_loss)
print(
str(t) + ':' + str(loss_val) + '\t' + str(s_loss) + '\t' +
str(c_loss))
new_image = sess.run(img_var)
imageio.imwrite(output + '\\final.jpg',
pre_img.deprocess_image(new_image[0]))